diff --git a/.github/workflows/release-tokenspeed-kernel.yml b/.github/workflows/release-tokenspeed-kernel.yml index 73213ac29..9e815bb36 100644 --- a/.github/workflows/release-tokenspeed-kernel.yml +++ b/.github/workflows/release-tokenspeed-kernel.yml @@ -216,7 +216,7 @@ jobs: cd /workspace/tokenspeed-kernel/python "${PYBIN}/python" -m pip install --upgrade \ - pip "setuptools<82" wheel build twine packaging "apache-tvm-ffi>=0.1.5,<=0.1.11" + pip "setuptools<82" wheel build twine packaging "apache-tvm-ffi==0.1.12" "${PYBIN}/python" -m pip install --no-cache-dir "torch==2.11.0" --index-url "${TORCH_INDEX_URL}" rm -rf build *.egg-info @@ -258,7 +258,7 @@ jobs: - name: Build source distribution run: | python -m pip install --upgrade \ - build twine "setuptools<82" wheel packaging "apache-tvm-ffi>=0.1.5,<=0.1.11" + build twine "setuptools<82" wheel packaging "apache-tvm-ffi==0.1.12" cd tokenspeed-kernel/python TOKENSPEED_KERNEL_BACKEND=cuda \ TOKENSPEED_KERNEL_VERSION_DATE=${{ needs.metadata.outputs.version_date }} \ diff --git a/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py b/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py index 57b777d06..9a10b245f 100644 --- a/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py +++ b/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py @@ -42,6 +42,13 @@ from tokenspeed.runtime.utils.env import global_server_args_dict +def supports_triton_rsag() -> bool: + platform = current_platform() + # The current NVIDIA multimem implementation is qualified on Hopper and + # datacenter Blackwell. Its SM120/SM121 launch corrupts CUDA state. + return not (platform.is_nvidia and platform.arch_version.major == 12) + + class TritonRSAGBackend: """Backend using TritonRSAG for token-aware reduce_scatter / all_gather. @@ -55,6 +62,12 @@ def __init__(self, fallback: CommBackend): # (group_tuple, hidden_size) -> Triton RS/AG state self._instances = {} + @staticmethod + def _is_supported_platform() -> bool: + # Keep token-aware collectives on the existing NCCL uneven-token path + # when the Triton symmetric-memory implementation is unavailable. + return supports_triton_rsag() + def _get_or_create(self, group: Group, hidden_size: int): key = (group, hidden_size) if key in self._instances: @@ -76,6 +89,8 @@ def all_gather( group: Group, dim: int = 0, ) -> torch.Tensor: + if not self._is_supported_platform(): + return self._fallback.all_gather(tensor, group=group, dim=dim) if tensor.dim() != 2: return self._fallback.all_gather(tensor, group=group, dim=dim) @@ -109,6 +124,12 @@ def token_all_gather( group: Group, scattered_num_tokens: list[int], ) -> torch.Tensor: + if not self._is_supported_platform(): + return self._fallback.token_all_gather( + tensor, + group=group, + scattered_num_tokens=scattered_num_tokens, + ) state = self._get_or_create(group, tensor.size(-1)) return all_gather(state, tensor, token_list_in_group=scattered_num_tokens) @@ -118,6 +139,12 @@ def token_reduce_scatter( group: Group, scattered_num_tokens: list[int], ) -> torch.Tensor: + if not self._is_supported_platform(): + return self._fallback.token_reduce_scatter( + tensor, + group=group, + scattered_num_tokens=scattered_num_tokens, + ) state = self._get_or_create(group, tensor.size(-1)) return reduce_scatter(state, tensor, token_list_in_group=scattered_num_tokens) diff --git a/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py b/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py index a8acb791d..0188e6166 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py +++ b/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py @@ -13,15 +13,13 @@ from __future__ import annotations +import tokenspeed_kernel import torch -from tokenspeed_kernel.ops.attention.flash_mla import ( - flash_mla_sparse_fwd, - flash_mla_with_kvcache, - get_mla_metadata, -) +from tokenspeed_kernel.ops.attention.flash_mla import flash_mla_sparse_fwd from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( deepseek_v4_indexer_decode_metadata_compute, ) +from tokenspeed_kernel.platform import Platform from tokenspeed_kernel.registry import error_fn try: @@ -62,6 +60,11 @@ DEEPSEEK_V4_DEFAULT_PREFILL_CHUNK_SIZE = 4 +def _use_flashinfer_sm120_sparse_mla() -> bool: + platform = Platform.get() + return platform.is_nvidia and platform.arch_version.major == 12 + + def _compressed_block_table_base_offsets( metadata: DeepseekV4ForwardMetadata, compress_ratio: int, @@ -265,7 +268,6 @@ def __init__(self, config) -> None: self.forward_metadata: DeepseekV4ForwardMetadata | None = None self.forward_prefill_metadata: DeepseekV4ForwardMetadata | None = None self.forward_decode_metadata: DeepseekV4ForwardMetadata | None = None - self._decode_tile_metadata = {} self._cuda_graph_metadata = {} self._cuda_graph_paged_cache_block_tables: dict[str, torch.Tensor] = {} # Per-sliding-group [max_bs] int32 buffers mirroring the block-table @@ -733,7 +735,6 @@ def init_forward_metadata( ) elif forward_mode is not None and forward_mode.is_extend_or_mixed(): self.forward_prefill_metadata = self.forward_metadata - self._decode_tile_metadata = {} def _update_decode_swa_metadata( self, @@ -947,24 +948,6 @@ def _dense_prefill_local_compressed_indices( out=out, ) - def _get_decode_tile_metadata(self, kind: str, bs: int): - phase = ( - "graph" - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() - else "eager" - ) - tile_metadata = self._decode_tile_metadata.get((phase, kind, bs)) - if tile_metadata is not None: - return tile_metadata - if get_mla_metadata is error_fn: - raise RuntimeError( - "DeepSeek V4 decode requires FlashMLA latent attention. " - "Build/install `tokenspeed-kernel/python` with FlashMLA." - ) - tile_metadata = get_mla_metadata()[0] - self._decode_tile_metadata[(phase, kind, bs)] = tile_metadata - return tile_metadata - def _fp8_ds_mla_cache_view( self, cache_2d: torch.Tensor, @@ -1020,12 +1003,6 @@ def forward_deepseek_v4_decode( f"metadata_tokens={metadata.token_to_req_indices.numel()}, " f"q_tokens={q.shape[0]}" ) - if flash_mla_with_kvcache is error_fn: - raise RuntimeError( - "DeepSeek V4 decode requires FlashMLA latent attention. " - "Build/install `tokenspeed-kernel/python` with FlashMLA." - ) - if q.shape[1] == padded_heads: q_padded = q.contiguous() else: @@ -1071,27 +1048,17 @@ def forward_deepseek_v4_decode( compressed_block_size, ) - out, _ = flash_mla_with_kvcache( - q=q_padded.unsqueeze(1), - k_cache=swa_cache, - block_table=None, - cache_seqlens=None, - head_dim_v=head_dim, - tile_scheduler_metadata=self._get_decode_tile_metadata( - kind, - q_padded.shape[0], - ), + out = tokenspeed_kernel.dsv4_sparse_mla_decode( + q=q_padded, + swa_kv_cache=swa_cache, + swa_indices=swa_indices, + swa_topk_lens=swa_lens, + compressed_kv_cache=compressed_cache, + compressed_indices=extra_indices, + compressed_topk_lens=extra_lens, softmax_scale=softmax_scale, - is_fp8_kvcache=True, - indices=swa_indices.unsqueeze(1), - attn_sink=attn_sink, - extra_k_cache=compressed_cache, - extra_indices_in_kvcache=extra_indices, - topk_length=swa_lens, - extra_topk_length=extra_lens, - ) - if out.dim() == 4: - out = out.squeeze(1) + sinks=attn_sink, + ) return out[:, :num_local_heads] def forward_deepseek_v4_mixed( @@ -1462,11 +1429,6 @@ def _forward_deepseek_v4_prefill_chunk( metadata = self.forward_metadata if metadata is None: raise RuntimeError("DeepSeek V4 prefill requires forward metadata") - if flash_mla_sparse_fwd is error_fn: - raise RuntimeError( - "DeepSeek V4 prefill requires FlashMLA sparse attention. " - "Build/install `tokenspeed-kernel/python` with FlashMLA." - ) with nvtx_range(f"attn_{kind}_prefill_pad_q"): if q.shape[1] == padded_heads: @@ -1478,6 +1440,55 @@ def _forward_deepseek_v4_prefill_chunk( device=q.device, ) q_padded[:, : q.shape[1]].copy_(q) + + if _use_flashinfer_sm120_sparse_mla(): + with nvtx_range(f"attn_{kind}_prefill_sparse_metadata"): + swa_block_size = token_to_kv_pool.swa_block_size + swa_indices, swa_lens = self._update_decode_swa_metadata( + metadata, + window_size=window_size, + block_size=swa_block_size, + ) + compressed_block_size = token_to_kv_pool.get_compressed_block_size( + layer_id + ) + extra_indices, extra_lens = ( + self._decode_compressed_attention_indices_and_lens( + positions, + compress_ratio=compress_ratio, + block_size=compressed_block_size, + topk_indices=topk_indices, + ) + ) + swa_cache = self._fp8_ds_mla_cache_view( + token_to_kv_pool.get_swa_kv_buffer(layer_id), + swa_block_size, + ) + compressed_cache = None + if compress_ratio > 1: + compressed_cache = self._fp8_ds_mla_cache_view( + token_to_kv_pool.get_compressed_kv_buffer_2d(layer_id), + compressed_block_size, + ) + with nvtx_range(f"attn_{kind}_prefill_flashinfer"): + out = tokenspeed_kernel.dsv4_sparse_mla_decode( + q=q_padded, + swa_kv_cache=swa_cache, + swa_indices=swa_indices, + swa_topk_lens=swa_lens, + compressed_kv_cache=compressed_cache, + compressed_indices=extra_indices, + compressed_topk_lens=extra_lens, + softmax_scale=softmax_scale, + sinks=attn_sink, + ) + return out[:, :num_local_heads] + + if flash_mla_sparse_fwd is error_fn: + raise RuntimeError( + "DeepSeek V4 prefill requires FlashMLA sparse attention. " + "Build/install `tokenspeed-kernel/python` with FlashMLA." + ) with nvtx_range(f"attn_{kind}_prefill_workspace"): kv_workspace, indices, lens = self._prefill_workspace( positions=positions, @@ -1612,7 +1623,6 @@ def init_cuda_graph_state( overlap_schedule_depth: int = 0, ): del seq_lens_buf - self._decode_tile_metadata = {} self._cuda_graph_max_tokens_per_req = max( 1, int(max_tokens_per_req), @@ -2086,7 +2096,6 @@ def advance_draft_forward_metadata(self, seq_lens: torch.Tensor | None = None): _refresh_decode_indexer_schedule_metadata(metadata) self.forward_decode_metadata = metadata self.forward_metadata = metadata - self._decode_tile_metadata = {} def forward_decode(self, *args, **kwargs): raise NotImplementedError("DeepSeek V4 uses the model-local attention forward") diff --git a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py index f2c3f8529..e8fe24db0 100644 --- a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py +++ b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py @@ -4,305 +4,8 @@ from __future__ import annotations -from functools import cache - +import tokenspeed_kernel import torch -import triton -import triton.language as tl - -from tokenspeed.runtime.utils import ceil_div - -try: - from tokenspeed_kernel.thirdparty import deep_gemm -except Exception: - deep_gemm = None # type: ignore[assignment] - - -@cache -def _compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: - device_props = torch.cuda.get_device_properties(0) - split_k = device_props.multi_processor_count // grid_size - if k is not None: - num_block_k = ceil_div(k, block_k) - split_k = min(split_k, num_block_k // 4) - return max(split_k, 1) - - -@triton.jit -def _load_reduced_mix( - gemm_out_mul, - token_id, - mix_id: tl.constexpr, - num_tokens, - hc_mult3: tl.constexpr, - n_splits: tl.constexpr, -): - value = tl.full((), 0.0, tl.float32) - for split_id in tl.static_range(0, n_splits): - offset = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 + mix_id - value += tl.load(gemm_out_mul + offset) - return value - - -@triton.jit -def _mhc_pre_mix_triton_kernel( - gemm_out_mul, - gemm_out_sqrsum, - hc_scale, - hc_base, - pre_mix, - post_mix, - comb_mix, - hidden_size: tl.constexpr, - rms_eps: tl.constexpr, - hc_eps: tl.constexpr, - sinkhorn_iters: tl.constexpr, - n_splits: tl.constexpr, - hc_mult: tl.constexpr, - hc_mult2: tl.constexpr, - hc_mult3: tl.constexpr, - block_comb: tl.constexpr, - num_tokens, -): - token_id = tl.program_id(0) - - rms_sum = tl.full((), 0.0, tl.float32) - for split_id in tl.static_range(0, n_splits): - rms_sum += tl.load(gemm_out_sqrsum + split_id * num_tokens + token_id) - rms = tl.rsqrt(rms_sum / (hc_mult * hidden_size) + rms_eps) - - pre_scale = tl.load(hc_scale) - for hc_id in tl.static_range(0, hc_mult): - mix = _load_reduced_mix( - gemm_out_mul, - token_id, - hc_id, - num_tokens, - hc_mult3, - n_splits, - ) - pre = tl.sigmoid(mix * rms * pre_scale + tl.load(hc_base + hc_id)) + hc_eps - tl.store(pre_mix + token_id * hc_mult + hc_id, pre) - - post_scale = tl.load(hc_scale + 1) - for hc_id in tl.static_range(0, hc_mult): - mix = _load_reduced_mix( - gemm_out_mul, - token_id, - hc_mult + hc_id, - num_tokens, - hc_mult3, - n_splits, - ) - post = ( - tl.sigmoid(mix * rms * post_scale + tl.load(hc_base + hc_mult + hc_id)) - * 2.0 - ) - tl.store(post_mix + token_id * hc_mult + hc_id, post) - - comb_offsets = tl.arange(0, block_comb) - comb_mask = comb_offsets < hc_mult2 - comb_scale = tl.load(hc_scale + 2) - comb_mix_values = tl.zeros((block_comb,), tl.float32) - for split_id in tl.static_range(0, n_splits): - split_base = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 - comb_mix_values += tl.load( - gemm_out_mul + split_base + hc_mult * 2 + comb_offsets, - mask=comb_mask, - other=0.0, - ) - comb_values = comb_mix_values * rms * comb_scale + tl.load( - hc_base + hc_mult * 2 + comb_offsets, mask=comb_mask, other=0.0 - ) - rows = comb_offsets // hc_mult - cols = comb_offsets - rows * hc_mult - active = comb_mask - - for row_id in tl.static_range(0, hc_mult): - row_values = tl.where((rows == row_id) & active, comb_values, -float("inf")) - row_max = tl.max(row_values, axis=0) - comb_values = tl.where( - (rows == row_id) & active, tl.exp(comb_values - row_max), comb_values - ) - for row_id in tl.static_range(0, hc_mult): - row_sum = tl.sum(tl.where((rows == row_id) & active, comb_values, 0.0), axis=0) - comb_values = tl.where( - (rows == row_id) & active, comb_values / row_sum + hc_eps, comb_values - ) - for col_id in tl.static_range(0, hc_mult): - col_sum = tl.sum(tl.where((cols == col_id) & active, comb_values, 0.0), axis=0) - comb_values = tl.where( - (cols == col_id) & active, - comb_values / (col_sum + hc_eps), - comb_values, - ) - - for _ in tl.static_range(1, sinkhorn_iters): - for row_id in tl.static_range(0, hc_mult): - row_sum = tl.sum( - tl.where((rows == row_id) & active, comb_values, 0.0), axis=0 - ) - comb_values = tl.where( - (rows == row_id) & active, - comb_values / (row_sum + hc_eps), - comb_values, - ) - for col_id in tl.static_range(0, hc_mult): - col_sum = tl.sum( - tl.where((cols == col_id) & active, comb_values, 0.0), axis=0 - ) - comb_values = tl.where( - (cols == col_id) & active, - comb_values / (col_sum + hc_eps), - comb_values, - ) - - tl.store( - comb_mix + token_id * hc_mult2 + comb_offsets, - comb_values, - mask=comb_mask, - ) - - -@triton.jit -def _mhc_pre_layer_triton_kernel( - pre_mix, - residual, - layer_input, - hidden_size: tl.constexpr, - hc_mult: tl.constexpr, - block_h: tl.constexpr, -): - token_id = tl.program_id(0) - hidden_block_id = tl.program_id(1) - - hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) - hidden_mask = hidden_offsets < hidden_size - layer_acc = tl.zeros((block_h,), tl.float32) - for hc_id in tl.static_range(0, hc_mult): - pre = tl.load(pre_mix + token_id * hc_mult + hc_id).to(tl.float32) - residual_offsets = ( - token_id * hc_mult * hidden_size + hc_id * hidden_size + hidden_offsets - ) - residual_values = tl.load( - residual + residual_offsets, mask=hidden_mask, other=0.0 - ).to(tl.float32) - layer_acc += pre * residual_values - tl.store( - layer_input + token_id * hidden_size + hidden_offsets, - layer_acc, - mask=hidden_mask, - ) - - -@triton.jit -def _mhc_post_triton_kernel( - comb, - residual, - post, - hidden_states, - out, - hidden_size: tl.constexpr, - hc_mult: tl.constexpr, - block_h: tl.constexpr, -): - token_id = tl.program_id(0) - hidden_block_id = tl.program_id(1) - hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) - hidden_mask = hidden_offsets < hidden_size - hidden_values = tl.load( - hidden_states + token_id * hidden_size + hidden_offsets, - mask=hidden_mask, - other=0.0, - ).to(tl.float32) - - for out_hc in tl.static_range(0, hc_mult): - acc = tl.load(post + token_id * hc_mult + out_hc).to(tl.float32) * hidden_values - for in_hc in tl.static_range(0, hc_mult): - comb_value = tl.load( - comb + token_id * hc_mult * hc_mult + in_hc * hc_mult + out_hc - ).to(tl.float32) - residual_values = tl.load( - residual - + token_id * hc_mult * hidden_size - + in_hc * hidden_size - + hidden_offsets, - mask=hidden_mask, - other=0.0, - ).to(tl.float32) - acc += comb_value * residual_values - tl.store( - out - + token_id * hc_mult * hidden_size - + out_hc * hidden_size - + hidden_offsets, - acc, - mask=hidden_mask, - ) - - -@triton.jit -def _mhc_post_hc4_triton_kernel( - comb, - residual, - post, - hidden_states, - out, - hidden_size: tl.constexpr, - block_h: tl.constexpr, -): - token_id = tl.program_id(0) - hidden_block_id = tl.program_id(1) - hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) - hidden_mask = hidden_offsets < hidden_size - token_hidden_offset = token_id * hidden_size - token_residual_offset = token_id * 4 * hidden_size - - hidden_values = tl.load( - hidden_states + token_hidden_offset + hidden_offsets, - mask=hidden_mask, - other=0.0, - ).to(tl.float32) - - post_base = token_id * 4 - acc0 = tl.load(post + post_base + 0).to(tl.float32) * hidden_values - acc1 = tl.load(post + post_base + 1).to(tl.float32) * hidden_values - acc2 = tl.load(post + post_base + 2).to(tl.float32) * hidden_values - acc3 = tl.load(post + post_base + 3).to(tl.float32) * hidden_values - - comb_base = token_id * 16 - for in_hc in tl.static_range(0, 4): - residual_values = tl.load( - residual + token_residual_offset + in_hc * hidden_size + hidden_offsets, - mask=hidden_mask, - other=0.0, - ).to(tl.float32) - comb_row = comb_base + in_hc * 4 - acc0 += tl.load(comb + comb_row + 0).to(tl.float32) * residual_values - acc1 += tl.load(comb + comb_row + 1).to(tl.float32) * residual_values - acc2 += tl.load(comb + comb_row + 2).to(tl.float32) * residual_values - acc3 += tl.load(comb + comb_row + 3).to(tl.float32) * residual_values - - tl.store( - out + token_residual_offset + hidden_offsets, - acc0, - mask=hidden_mask, - ) - tl.store( - out + token_residual_offset + hidden_size + hidden_offsets, - acc1, - mask=hidden_mask, - ) - tl.store( - out + token_residual_offset + hidden_size * 2 + hidden_offsets, - acc2, - mask=hidden_mask, - ) - tl.store( - out + token_residual_offset + hidden_size * 3 + hidden_offsets, - acc3, - mask=hidden_mask, - ) def mhc_fused_hc( @@ -317,10 +20,7 @@ def mhc_fused_hc( hc_eps: float, sinkhorn_iters: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Fused post_mapping(prev) + pre_mapping(curr). - - Returns (residual_cur, layer_input, post_cur, comb_cur). - """ + """Apply the previous post-map followed by the current pre-map.""" residual_cur = mhc_post(x_prev, residual_prev, post_prev, comb_prev) layer_input, post_cur, comb_cur = mhc_pre( residual_cur, @@ -343,109 +43,15 @@ def mhc_pre( hc_eps: float, sinkhorn_iters: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if residual.dtype != torch.bfloat16 or fn.dtype != torch.float32: - raise RuntimeError("fast mHC requires bf16 residual and fp32 weights") - if not residual.is_cuda: - raise RuntimeError("fast mHC requires CUDA tensors") - - if deep_gemm is None: - raise RuntimeError("deep_gemm.tf32_hc_prenorm_gemm is unavailable") - - hc_mult = residual.shape[-2] - hidden_size = residual.shape[-1] - hc_mult2 = hc_mult * hc_mult - hc_mult3 = hc_mult * 2 + hc_mult2 - hc_hidden_size = hc_mult * hidden_size - outer_shape = residual.shape[:-2] - residual_flat = residual.view(-1, hc_mult, hidden_size) - num_tokens = residual_flat.shape[0] - if num_tokens == 0: - return ( - residual.new_empty(*outer_shape, hidden_size), - torch.empty( - *outer_shape, - hc_mult, - 1, - dtype=torch.float32, - device=residual.device, - ), - torch.empty( - *outer_shape, - hc_mult, - hc_mult, - dtype=torch.float32, - device=residual.device, - ), - ) - - block_k = 64 - block_m = 64 - n_splits = _compute_num_split( - block_k, hc_hidden_size, ceil_div(num_tokens, block_m) - ) - - post_mix = torch.empty( - num_tokens, hc_mult, dtype=torch.float32, device=residual.device - ) - pre_mix = torch.empty( - num_tokens, hc_mult, dtype=torch.float32, device=residual.device - ) - comb_mix = torch.empty( - num_tokens, hc_mult2, dtype=torch.float32, device=residual.device - ) - layer_input = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device - ) - gemm_out_mul = torch.empty( - n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device - ) - gemm_out_sqrsum = torch.empty( - n_splits, num_tokens, dtype=torch.float32, device=residual.device - ) - - deep_gemm.tf32_hc_prenorm_gemm( - residual_flat.view(num_tokens, hc_hidden_size), + """Delegate the DSv4 pre-map to the tokenspeed-kernel boundary.""" + return tokenspeed_kernel.mhc_pre( + residual, fn, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, - ) - block_h = 1024 - block_comb = triton.next_power_of_2(hc_mult2) - _mhc_pre_mix_triton_kernel[(num_tokens,)]( - gemm_out_mul, - gemm_out_sqrsum, hc_scale, hc_base, - pre_mix, - post_mix, - comb_mix, - hidden_size=hidden_size, - rms_eps=rms_eps, - hc_eps=hc_eps, - sinkhorn_iters=sinkhorn_iters, - n_splits=n_splits, - hc_mult=hc_mult, - hc_mult2=hc_mult2, - hc_mult3=hc_mult3, - block_comb=block_comb, - num_tokens=num_tokens, - num_warps=1, - ) - _mhc_pre_layer_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( - pre_mix, - residual_flat, - layer_input, - hidden_size=hidden_size, - hc_mult=hc_mult, - block_h=block_h, - num_warps=4, - ) - - return ( - layer_input.view(*outer_shape, hidden_size), - post_mix.view(*outer_shape, hc_mult, 1), - comb_mix.view(*outer_shape, hc_mult, hc_mult), + rms_eps, + hc_eps, + sinkhorn_iters, ) @@ -455,42 +61,5 @@ def mhc_post( post: torch.Tensor, comb: torch.Tensor, ) -> torch.Tensor: - if not hidden_states.is_cuda: - raise RuntimeError("fast mHC requires CUDA tensors") - if residual.numel() == 0: - return torch.empty_like(residual) - out = torch.empty_like(residual) - hc_mult = residual.shape[-2] - hidden_size = residual.shape[-1] - residual_flat = residual.view(-1, hc_mult, hidden_size) - hidden_states_flat = hidden_states.view(-1, hidden_size) - post_flat = post.view(-1, hc_mult) - comb_flat = comb.view(-1, hc_mult, hc_mult) - num_tokens = residual_flat.shape[0] - if hc_mult == 4: - block_h = 256 - _mhc_post_hc4_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( - comb_flat, - residual_flat, - post_flat, - hidden_states_flat, - out, - hidden_size=hidden_size, - block_h=block_h, - num_warps=4, - ) - return out - - block_h = 1024 - _mhc_post_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( - comb_flat, - residual_flat, - post_flat, - hidden_states_flat, - out, - hidden_size=hidden_size, - hc_mult=hc_mult, - block_h=block_h, - num_warps=4, - ) - return out + """Delegate the DSv4 post-map to the tokenspeed-kernel boundary.""" + return tokenspeed_kernel.mhc_post(hidden_states, residual, post, comb) diff --git a/python/tokenspeed/runtime/layers/logits_processor.py b/python/tokenspeed/runtime/layers/logits_processor.py index 7e63df6cb..7d1c85158 100755 --- a/python/tokenspeed/runtime/layers/logits_processor.py +++ b/python/tokenspeed/runtime/layers/logits_processor.py @@ -34,6 +34,9 @@ from tokenspeed_kernel.platform import current_platform from torch import nn +from tokenspeed.runtime.distributed.comm_backend.triton_rsag import ( + supports_triton_rsag, +) from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor from tokenspeed.runtime.distributed.process_group_manager import ( process_group_manager as pg_manager, @@ -286,7 +289,7 @@ def _resolve_logits_layout_plan( ) def _init_all_gather_state(self, lm_head: VocabParallelEmbedding): - if not current_platform().is_nvidia: + if not current_platform().is_nvidia or not supports_triton_rsag(): return None if self.tp_size == 1 or self.skip_all_gather: @@ -307,7 +310,7 @@ def _init_all_gather_state(self, lm_head: VocabParallelEmbedding): return self._LOGITS_AG_STATES[key] def _init_dist_argmax_state(self, lm_head: VocabParallelEmbedding): - if not current_platform().is_nvidia: + if not current_platform().is_nvidia or not supports_triton_rsag(): return None if self.tp_size == 1 or self.skip_all_gather or self.dp_sampling_enabled: diff --git a/python/tokenspeed/runtime/models/deepseek_v4.py b/python/tokenspeed/runtime/models/deepseek_v4.py index 18b09b5dd..fdde90453 100644 --- a/python/tokenspeed/runtime/models/deepseek_v4.py +++ b/python/tokenspeed/runtime/models/deepseek_v4.py @@ -34,6 +34,7 @@ import torch import torch.nn.functional as F +from tokenspeed_kernel import prebuild_dsv4_sparse_mla try: # Optional dependency; the module-level wrapper imports the external @@ -45,8 +46,10 @@ from tokenspeed_kernel.ops.attention.cuda.deepseek_v4 import ( has_indexer_mxfp4_paged_gather, + has_indexer_topk_prefill, has_persistent_topk, indexer_mxfp4_paged_gather, + indexer_topk_prefill, persistent_topk, ) from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( @@ -712,6 +715,44 @@ def _deepseek_v4_indexer_topk_from_logits( return topk +def _deepseek_v4_launch_indexer_topk_prefill( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + output: torch.Tensor, + topk_tokens: int, +) -> None: + """Launch the architecture-compatible DeepSeek V4 prefill selector.""" + + if _platform.is_nvidia and _platform.arch_version.major == 12: + if not has_indexer_topk_prefill(): + raise RuntimeError( + "DeepSeek V4 prefill indexer on SM120 requires the " + "TokenSpeed CUDA prefill top-k op" + ) + indexer_topk_prefill( + logits, + row_starts, + row_ends, + output, + topk_tokens, + ) + return + + trtllm_ops = getattr(torch.ops, "trtllm", None) + if trtllm_ops is None or not hasattr(trtllm_ops, "indexer_topk_prefill"): + raise RuntimeError( + "DeepSeek V4 prefill indexer requires the CUDA prefill top-k op" + ) + trtllm_ops.indexer_topk_prefill( + logits, + row_starts, + row_ends, + output, + topk_tokens, + ) + + def _deepseek_v4_indexer_topk_from_logits_prefill_op( logits: torch.Tensor, length_rows: torch.Tensor, @@ -721,15 +762,10 @@ def _deepseek_v4_indexer_topk_from_logits_prefill_op( row_ends: torch.Tensor | None = None, out: torch.Tensor, ) -> torch.Tensor: - """Use the local TRT-LLM CUDA prefill selector.""" + """Use the architecture-compatible CUDA prefill selector.""" if not logits.is_cuda or logits.dtype != torch.float32: raise RuntimeError("DeepSeek V4 prefill indexer requires CUDA float32 logits") - trtllm_ops = getattr(torch.ops, "trtllm", None) - if trtllm_ops is None or not hasattr(trtllm_ops, "indexer_topk_prefill"): - raise RuntimeError( - "DeepSeek V4 prefill indexer requires the CUDA prefill top-k op" - ) num_rows = length_rows.numel() if num_rows == 0: @@ -767,7 +803,7 @@ def _deepseek_v4_indexer_topk_from_logits_prefill_op( topk = out[:num_rows] topk.fill_(-1) - trtllm_ops.indexer_topk_prefill( + _deepseek_v4_launch_indexer_topk_prefill( logits, row_starts_for_kernel, row_ends_for_kernel, @@ -1379,6 +1415,7 @@ def _deepseek_v4_indexer_topk_prefill_deepgemm( use_prefill_topk_op: bool, gathered_k: tuple[torch.Tensor, torch.Tensor] | None = None, gather_workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: q_values, q_scales = index_q if not _deepseek_v4_deepgemm_fp4_indexer_available(q_values): @@ -1429,15 +1466,14 @@ def _deepseek_v4_indexer_topk_prefill_deepgemm( ) with nvtx_range("indexer_topk_prefill_select"): - return ( - _deepseek_v4_indexer_topk_from_logits( - logits, - row_lens, - topk_tokens, - use_prefill_topk_op=use_prefill_topk_op, - ), - gathered_k, + topk = _deepseek_v4_indexer_topk_from_logits( + logits, + row_lens, + topk_tokens, + use_prefill_topk_op=use_prefill_topk_op, + out=out, ) + return topk, gathered_k def _deepseek_v4_indexer_topk_from_cache_deepgemm_decode( @@ -1640,6 +1676,11 @@ def fill_prefill() -> None: raise RuntimeError( "DeepSeek V4 sparse indexer prefill metadata is incomplete" ) + direct_out = ( + topk_out[token_start:token_end] + if _platform.is_nvidia and _platform.arch_version.major == 12 + else None + ) topk, next_gathered_k = _deepseek_v4_indexer_topk_prefill_deepgemm( cache_2d=cache_2d, block_table=block_table[req_start:req_end], @@ -1658,11 +1699,13 @@ def fill_prefill() -> None: use_prefill_topk_op=True, gathered_k=reuse_k, gather_workspace=gather_workspace, + out=direct_out, ) if next_gathered_k is not None: gather_cache_key = key gathered_k = next_gathered_k - topk_out[token_start:token_end].copy_(topk) + if direct_out is None: + topk_out[token_start:token_end].copy_(topk) def fill_decode() -> None: if num_decode_tokens <= 0: @@ -4187,6 +4230,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): del params_dict, moe_loader self.post_load_weights() self.warmup_deep_gemm() + self.warmup_flashinfer() def post_load_weights(self): mega_moe_experts: list[DeepseekV4MegaMoEExperts] = [] @@ -4199,6 +4243,11 @@ def post_load_weights(self): elif isinstance(module, MoELayer): module.process_weights_after_loading(module) + def warmup_flashinfer(self) -> None: + platform = current_platform() + if platform.is_nvidia and platform.arch_version.major == 12: + prebuild_dsv4_sparse_mla() + def warmup_deep_gemm(self) -> None: """Pre-compile all DeepGEMM JIT kernels used by this model. diff --git a/test/runtime/distributed/test_triton_rsag_platform.py b/test/runtime/distributed/test_triton_rsag_platform.py new file mode 100644 index 000000000..bd5cc877f --- /dev/null +++ b/test/runtime/distributed/test_triton_rsag_platform.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from tokenspeed_kernel.platform import ArchVersion + +import tokenspeed.runtime.distributed.comm_backend.triton_rsag as triton_rsag + + +class _Fallback: + def __init__(self) -> None: + self.calls: list[tuple[str, object]] = [] + + def token_all_gather(self, tensor, group, scattered_num_tokens): + self.calls.append(("token_all_gather", (tensor, group, scattered_num_tokens))) + return tensor + 1 + + def token_reduce_scatter(self, tensor, group, scattered_num_tokens): + self.calls.append( + ("token_reduce_scatter", (tensor, group, scattered_num_tokens)) + ) + return tensor + 2 + + def all_gather(self, tensor, group, dim=0): + self.calls.append(("all_gather", (tensor, group, dim))) + return tensor + 3 + + +def _platform(major: int, minor: int = 0) -> SimpleNamespace: + return SimpleNamespace( + is_nvidia=True, + arch_version=ArchVersion(major, minor), + ) + + +def test_sm120_falls_back_for_all_triton_rsag_entrypoints(monkeypatch) -> None: + fallback = _Fallback() + backend = triton_rsag.TritonRSAGBackend(fallback=fallback) + monkeypatch.setattr(triton_rsag, "current_platform", lambda: _platform(12)) + monkeypatch.setattr( + backend, + "_get_or_create", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("SM120 must not create Triton RSAG state") + ), + ) + tensor = torch.zeros(1, 8, dtype=torch.bfloat16) + group = (0, 1, 2, 3) + scattered = [1, 1, 1, 1] + + gathered = backend.token_all_gather(tensor, group, scattered) + scattered_out = backend.token_reduce_scatter(tensor, group, scattered) + inner = backend.all_gather(tensor, group, dim=-1) + + torch.testing.assert_close(gathered, tensor + 1) + torch.testing.assert_close(scattered_out, tensor + 2) + torch.testing.assert_close(inner, tensor + 3) + assert [name for name, _args in fallback.calls] == [ + "token_all_gather", + "token_reduce_scatter", + "all_gather", + ] + + +@pytest.mark.parametrize("major", [9, 10], ids=["h100", "b200"]) +def test_existing_nvidia_architectures_preserve_triton_rsag_token_path( + major: int, + monkeypatch, +) -> None: + fallback = _Fallback() + backend = triton_rsag.TritonRSAGBackend(fallback=fallback) + monkeypatch.setattr(triton_rsag, "current_platform", lambda: _platform(major)) + state = object() + monkeypatch.setattr(backend, "_get_or_create", lambda *_args: state) + expected = torch.ones(4, 8, dtype=torch.bfloat16) + call: dict[str, object] = {} + + def fake_all_gather(state_arg, tensor, token_list_in_group): + call.update( + state=state_arg, + tensor=tensor, + token_list_in_group=token_list_in_group, + ) + return expected + + monkeypatch.setattr(triton_rsag, "all_gather", fake_all_gather) + tensor = torch.zeros(1, 8, dtype=torch.bfloat16) + + actual = backend.token_all_gather(tensor, (0, 1, 2, 3), [1, 1, 1, 1]) + + assert actual is expected + assert call == { + "state": state, + "tensor": tensor, + "token_list_in_group": [1, 1, 1, 1], + } + assert fallback.calls == [] diff --git a/test/runtime/layers/test_deepseek_v4_mhc_boundary.py b/test/runtime/layers/test_deepseek_v4_mhc_boundary.py new file mode 100644 index 000000000..b65a144a7 --- /dev/null +++ b/test/runtime/layers/test_deepseek_v4_mhc_boundary.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +import tokenspeed_kernel +import torch + +from tokenspeed.runtime.layers import deepseek_v4_mhc + + +def test_runtime_mhc_calls_tokenspeed_kernel_boundary(monkeypatch) -> None: + x_prev = torch.randn(2, 8, dtype=torch.bfloat16) + residual_prev = torch.randn(2, 4, 8, dtype=torch.bfloat16) + post_prev = torch.randn(2, 4, 1, dtype=torch.float32) + comb_prev = torch.randn(2, 4, 4, dtype=torch.float32) + fn = torch.randn(24, 32, dtype=torch.float32) + scale = torch.ones(3, dtype=torch.float32) + base = torch.zeros(24, dtype=torch.float32) + residual_cur = torch.empty_like(residual_prev) + layer_input = torch.empty_like(x_prev) + post_cur = torch.empty_like(post_prev) + comb_cur = torch.empty_like(comb_prev) + calls: list[tuple[str, tuple[object, ...]]] = [] + + def fake_post(*args): + calls.append(("post", args)) + return residual_cur + + def fake_pre(*args): + calls.append(("pre", args)) + return layer_input, post_cur, comb_cur + + monkeypatch.setattr(tokenspeed_kernel, "mhc_post", fake_post) + monkeypatch.setattr(tokenspeed_kernel, "mhc_pre", fake_pre) + + result = deepseek_v4_mhc.mhc_fused_hc( + x_prev, + residual_prev, + post_prev, + comb_prev, + fn, + scale, + base, + 1e-6, + 2e-6, + 3, + ) + + assert result == (residual_cur, layer_input, post_cur, comb_cur) + assert calls == [ + ("post", (x_prev, residual_prev, post_prev, comb_prev)), + ("pre", (residual_cur, fn, scale, base, 1e-6, 2e-6, 3)), + ] diff --git a/test/runtime/layers/test_deepseek_v4_sparse_mla_boundary.py b/test/runtime/layers/test_deepseek_v4_sparse_mla_boundary.py new file mode 100644 index 000000000..d651ad056 --- /dev/null +++ b/test/runtime/layers/test_deepseek_v4_sparse_mla_boundary.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +from types import SimpleNamespace + +import tokenspeed_kernel +import torch + +import tokenspeed.runtime.layers.attention.backends.deepseek_v4 as deepseek_v4_backend +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.deepseek_v4 import ( + DeepseekV4AttentionBackend, +) + + +def test_deepseek_v4_decode_calls_tokenspeed_kernel_sparse_mla_boundary( + monkeypatch, +) -> None: + tokens, local_heads, padded_heads, head_dim = 2, 32, 64, 512 + swa_indices = torch.zeros(tokens, 128, dtype=torch.int32) + swa_lens = torch.full((tokens,), 128, dtype=torch.int32) + compressed_indices = torch.zeros(tokens, 1, 512, dtype=torch.int32) + compressed_lens = torch.full((tokens,), 512, dtype=torch.int32) + attention = SimpleNamespace( + decode_swa_indices=swa_indices, + decode_swa_lens=swa_lens, + decode_swa_window_size=128, + decode_swa_block_size=64, + ) + metadata = SimpleNamespace( + forward_mode=ForwardMode.DECODE, + token_to_req_indices=torch.arange(tokens, dtype=torch.int32), + attention=attention, + ) + backend = object.__new__(DeepseekV4AttentionBackend) + backend.forward_metadata = metadata + monkeypatch.setattr(backend, "_select_decode_metadata", lambda _tokens: metadata) + monkeypatch.setattr( + backend, + "_decode_compressed_attention_indices_and_lens", + lambda *_args, **_kwargs: (compressed_indices, compressed_lens), + ) + + swa_cache = torch.empty(3, 64, 1, 584, dtype=torch.uint8) + compressed_cache = torch.empty(4, 2, 1, 584, dtype=torch.uint8) + monkeypatch.setattr( + backend, + "_fp8_ds_mla_cache_view", + lambda cache, _block_size: cache, + ) + pool = SimpleNamespace( + swa_block_size=64, + get_compressed_block_size=lambda _layer_id: 2, + get_swa_kv_buffer=lambda _layer_id: swa_cache, + get_compressed_kv_buffer_2d=lambda _layer_id: compressed_cache, + ) + expected = torch.randn(tokens, padded_heads, head_dim, dtype=torch.bfloat16) + call: dict[str, object] = {} + + def fake_sparse_mla_decode(**kwargs): + call.update(kwargs) + return expected + + monkeypatch.setattr( + tokenspeed_kernel, + "dsv4_sparse_mla_decode", + fake_sparse_mla_decode, + ) + q = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16) + positions = torch.arange(tokens, dtype=torch.int64) + sinks = torch.zeros(padded_heads, dtype=torch.float32) + + actual = backend.forward_deepseek_v4_decode( + q=q, + positions=positions, + token_to_kv_pool=pool, + layer_id=0, + kind="csa", + compress_ratio=4, + num_local_heads=local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=128, + softmax_scale=head_dim**-0.5, + attn_sink=sinks, + topk_indices=torch.zeros(tokens, 512, dtype=torch.int32), + ) + + torch.testing.assert_close(actual, expected[:, :local_heads]) + q_padded = call["q"] + assert isinstance(q_padded, torch.Tensor) + assert q_padded.shape == (tokens, padded_heads, head_dim) + torch.testing.assert_close(q_padded[:, :local_heads], q) + assert call == { + "q": q_padded, + "swa_kv_cache": swa_cache, + "swa_indices": swa_indices, + "swa_topk_lens": swa_lens, + "compressed_kv_cache": compressed_cache, + "compressed_indices": compressed_indices, + "compressed_topk_lens": compressed_lens, + "softmax_scale": head_dim**-0.5, + "sinks": sinks, + } + + +def test_deepseek_v4_sm120_prefill_calls_tokenspeed_kernel_sparse_mla_boundary( + monkeypatch, +) -> None: + tokens, local_heads, padded_heads, head_dim = 128, 32, 64, 512 + swa_indices = torch.zeros(tokens, 128, dtype=torch.int32) + swa_lens = torch.full((tokens,), 128, dtype=torch.int32) + compressed_indices = torch.zeros(tokens, 1, 512, dtype=torch.int32) + compressed_lens = torch.full((tokens,), 512, dtype=torch.int32) + backend = object.__new__(DeepseekV4AttentionBackend) + backend.forward_metadata = SimpleNamespace() + monkeypatch.setattr( + deepseek_v4_backend, + "_use_flashinfer_sm120_sparse_mla", + lambda: True, + raising=False, + ) + monkeypatch.setattr( + backend, + "_update_decode_swa_metadata", + lambda *_args, **_kwargs: (swa_indices, swa_lens), + ) + monkeypatch.setattr( + backend, + "_decode_compressed_attention_indices_and_lens", + lambda *_args, **_kwargs: (compressed_indices, compressed_lens), + ) + monkeypatch.setattr( + backend, + "_prefill_workspace", + lambda **_kwargs: (_ for _ in ()).throw( + AssertionError("SM120 prefill must not gather for FlashMLA") + ), + ) + + swa_cache = torch.empty(3, 64, 1, 584, dtype=torch.uint8) + compressed_cache = torch.empty(4, 2, 1, 584, dtype=torch.uint8) + monkeypatch.setattr( + backend, + "_fp8_ds_mla_cache_view", + lambda cache, _block_size: cache, + ) + pool = SimpleNamespace( + swa_block_size=64, + get_compressed_block_size=lambda _layer_id: 2, + get_swa_kv_buffer=lambda _layer_id: swa_cache, + get_compressed_kv_buffer_2d=lambda _layer_id: compressed_cache, + ) + expected = torch.randn(tokens, padded_heads, head_dim, dtype=torch.bfloat16) + call: dict[str, object] = {} + + def fake_sparse_mla(**kwargs): + call.update(kwargs) + return expected + + monkeypatch.setattr(tokenspeed_kernel, "dsv4_sparse_mla_decode", fake_sparse_mla) + q = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16) + positions = torch.arange(tokens, dtype=torch.int64) + sinks = torch.zeros(padded_heads, dtype=torch.float32) + + actual = backend._forward_deepseek_v4_prefill_chunk( + q=q, + positions=positions, + token_to_kv_pool=pool, + layer_id=0, + kind="csa", + compress_ratio=4, + num_local_heads=local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=128, + softmax_scale=head_dim**-0.5, + attn_sink=sinks, + topk_indices=torch.zeros(tokens, 512, dtype=torch.int32), + ) + + torch.testing.assert_close(actual, expected[:, :local_heads]) + q_padded = call["q"] + assert isinstance(q_padded, torch.Tensor) + assert q_padded.shape == (tokens, padded_heads, head_dim) + torch.testing.assert_close(q_padded[:, :local_heads], q) + assert call == { + "q": q_padded, + "swa_kv_cache": swa_cache, + "swa_indices": swa_indices, + "swa_topk_lens": swa_lens, + "compressed_kv_cache": compressed_cache, + "compressed_indices": compressed_indices, + "compressed_topk_lens": compressed_lens, + "softmax_scale": head_dim**-0.5, + "sinks": sinks, + } + + +def test_deepseek_v4_non_sm120_prefill_retains_flashmla_path(monkeypatch) -> None: + tokens, heads, head_dim = 2, 32, 512 + backend = object.__new__(DeepseekV4AttentionBackend) + backend.forward_metadata = SimpleNamespace() + monkeypatch.setattr( + deepseek_v4_backend, + "_use_flashinfer_sm120_sparse_mla", + lambda: False, + raising=False, + ) + kv_workspace = torch.zeros(1, 8, head_dim, dtype=torch.bfloat16) + indices = torch.zeros(tokens, 8, dtype=torch.int32) + lens = torch.full((tokens,), 8, dtype=torch.int32) + monkeypatch.setattr( + backend, + "_prefill_workspace", + lambda **_kwargs: (kv_workspace, indices, lens), + ) + expected = torch.randn(tokens, heads, head_dim, dtype=torch.bfloat16) + call: dict[str, object] = {} + + def fake_flashmla(**kwargs): + call.update(kwargs) + return expected, None, None + + monkeypatch.setattr(deepseek_v4_backend, "flash_mla_sparse_fwd", fake_flashmla) + monkeypatch.setattr( + tokenspeed_kernel, + "dsv4_sparse_mla_decode", + lambda **_kwargs: (_ for _ in ()).throw( + AssertionError("non-SM120 prefill must retain FlashMLA") + ), + ) + q = torch.randn(tokens, heads, head_dim, dtype=torch.bfloat16) + sinks = torch.zeros(heads, dtype=torch.float32) + + actual = backend._forward_deepseek_v4_prefill_chunk( + q=q, + positions=torch.arange(tokens, dtype=torch.int64), + token_to_kv_pool=SimpleNamespace(), + layer_id=0, + kind="csa", + compress_ratio=4, + num_local_heads=heads, + padded_heads=heads, + head_dim=head_dim, + window_size=128, + softmax_scale=head_dim**-0.5, + attn_sink=sinks, + topk_indices=torch.zeros(tokens, 8, dtype=torch.int32), + ) + + torch.testing.assert_close(actual, expected) + assert set(call) == { + "q", + "kv", + "indices", + "sm_scale", + "attn_sink", + "topk_length", + } + torch.testing.assert_close(call["q"], q) + torch.testing.assert_close(call["kv"], kv_workspace.view(-1, 1, head_dim)) + torch.testing.assert_close(call["indices"], indices.unsqueeze(1)) + assert call["sm_scale"] == head_dim**-0.5 + torch.testing.assert_close(call["attn_sink"], sinks) + torch.testing.assert_close(call["topk_length"], lens) diff --git a/test/runtime/models/test_deepseek_v4_sm120_indexer.py b/test/runtime/models/test_deepseek_v4_sm120_indexer.py new file mode 100644 index 000000000..dce292007 --- /dev/null +++ b/test/runtime/models/test_deepseek_v4_sm120_indexer.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +import tokenspeed.runtime.models.deepseek_v4 as deepseek_v4_model + + +@pytest.mark.parametrize( + ("arch_major", "expected_backend"), + [(12, "tokenspeed"), (10, "trtllm")], +) +def test_prefill_topk_uses_tokenspeed_cuda_only_on_sm120( + monkeypatch, + arch_major: int, + expected_backend: str, +) -> None: + monkeypatch.setattr( + deepseek_v4_model, + "_platform", + SimpleNamespace( + is_nvidia=True, + arch_version=SimpleNamespace(major=arch_major), + ), + ) + calls: list[str] = [] + + def fake_tokenspeed_topk(*args, **kwargs) -> None: + del args, kwargs + calls.append("tokenspeed") + + def fake_trtllm_topk(*args, **kwargs) -> None: + del args, kwargs + calls.append("trtllm") + + monkeypatch.setattr( + deepseek_v4_model, + "has_indexer_topk_prefill", + lambda: True, + raising=False, + ) + monkeypatch.setattr( + deepseek_v4_model, + "indexer_topk_prefill", + fake_tokenspeed_topk, + raising=False, + ) + logits = torch.empty((2, 4), dtype=torch.float32) + row_starts = torch.zeros(2, dtype=torch.int32) + row_ends = torch.full((2,), 4, dtype=torch.int32) + output = torch.empty((2, 2), dtype=torch.int32) + + with patch.object( + torch.ops.trtllm, + "indexer_topk_prefill", + fake_trtllm_topk, + create=True, + ): + deepseek_v4_model._deepseek_v4_launch_indexer_topk_prefill( + logits, + row_starts, + row_ends, + output, + 2, + ) + + assert calls == [expected_backend] + + +@pytest.mark.parametrize( + ("arch_major", "expects_direct_out"), + [(12, True), (10, False)], +) +def test_prefill_indexer_writes_directly_only_on_sm120( + monkeypatch, + arch_major: int, + expects_direct_out: bool, +) -> None: + monkeypatch.setattr( + deepseek_v4_model, + "_platform", + SimpleNamespace( + is_nvidia=True, + arch_version=SimpleNamespace(major=arch_major), + ), + ) + captured: dict[str, object] = {} + + def fake_prefill_topk(**kwargs): + out = kwargs.get("out") + captured["out"] = out + if out is None: + topk = torch.full((2, 4), 7, dtype=torch.int32) + else: + assert isinstance(out, torch.Tensor) + out.fill_(7) + topk = out + return topk, None + + monkeypatch.setattr( + deepseek_v4_model, + "_deepseek_v4_indexer_topk_prefill_deepgemm", + fake_prefill_topk, + ) + topk_buffer = torch.empty((2, 4), dtype=torch.int32) + + actual = deepseek_v4_model._deepseek_v4_sparse_attn_indexer_native( + cache_2d=torch.empty((1, 1), dtype=torch.uint8), + positions=torch.arange(2, dtype=torch.int64), + token_to_req_indices=torch.zeros(2, dtype=torch.int32), + block_table=torch.zeros((1, 1), dtype=torch.int32), + seq_lens_cpu=torch.tensor([2], dtype=torch.int32), + query_lens_cpu=torch.tensor([2], dtype=torch.int32), + prefill_chunk_specs=torch.tensor([[0, 2, 0, 1, 0]], dtype=torch.int64), + prefill_chunk_offsets=torch.tensor([[0, 0, 0, 1, 1, 0, 2]], dtype=torch.int64), + prefill_slots=torch.empty(0, dtype=torch.int64), + prefill_cu_seq_lens=torch.tensor([0, 1], dtype=torch.int32), + prefill_cu_seqlen_k_start=torch.tensor([0], dtype=torch.int32), + prefill_cu_seqlen_k_end=torch.tensor([1], dtype=torch.int32), + prefill_seq_lens_k=torch.tensor([1], dtype=torch.int32), + packed_q_values=torch.empty((2, 1, 1), dtype=torch.int8), + packed_q_scales=torch.empty((2, 1), dtype=torch.int32), + packed_weights=torch.empty((2, 1), dtype=torch.float32), + decode_schedule_metadata=None, + decode_context_lens=None, + decode_block_table=None, + decode_max_context_len=0, + topk_indices_buffer=topk_buffer, + prefill_gather_values_workspace=torch.empty((0, 1), dtype=torch.uint8), + prefill_gather_scales_workspace=torch.empty((0, 1), dtype=torch.uint8), + persistent_topk_workspace=torch.empty(0, dtype=torch.uint8), + cache_block_size=1, + compress_ratio=4, + topk_tokens=4, + num_prefill_tokens=2, + num_decode_tokens=0, + ) + + if expects_direct_out: + assert captured["out"] is not None + assert isinstance(captured["out"], torch.Tensor) + assert captured["out"].data_ptr() == topk_buffer.data_ptr() + else: + assert captured["out"] is None + torch.testing.assert_close(actual, torch.full_like(actual, 7)) diff --git a/test/runtime/test_logits_processor.py b/test/runtime/test_logits_processor.py index b5c4f7544..5ac5d64a3 100644 --- a/test/runtime/test_logits_processor.py +++ b/test/runtime/test_logits_processor.py @@ -90,6 +90,84 @@ def fake_all_gather_into_tensor(output, input_, group): assert tuple(output.next_token_logits.shape) == (0, 6) +def test_sm120_tp_logits_skips_symmetric_memory_states(monkeypatch): + processor = LogitsProcessor( + config=SimpleNamespace(model_type="test", vocab_size=8192), + tp_rank=0, + tp_size=2, + tp_group=(0, 1), + ) + lm_head = SimpleNamespace(weight=torch.ones((4096, 2), dtype=torch.bfloat16)) + monkeypatch.setattr( + logits_processor_module, + "current_platform", + lambda: SimpleNamespace( + is_nvidia=True, + arch_version=SimpleNamespace(major=12), + ), + ) + monkeypatch.setattr(logits_processor_module, "supports_triton_rsag", lambda: False) + monkeypatch.setattr( + logits_processor_module, + "create_state", + lambda *_args, **_kwargs: pytest.fail( + "SM120 logits gather must not create Triton RSAG state" + ), + ) + monkeypatch.setattr( + logits_processor_module, + "create_dist_argmax_state", + lambda *_args, **_kwargs: pytest.fail( + "SM120 logits argmax must not create symmetric-memory state" + ), + ) + + assert processor._init_all_gather_state(lm_head) is None + assert processor._init_dist_argmax_state(lm_head) is None + + +@pytest.mark.parametrize("major", [9, 10], ids=["h100", "b200"]) +def test_existing_nvidia_architectures_keep_tp_logits_fast_path(major, monkeypatch): + LogitsProcessor._LOGITS_AG_STATES.clear() + processor = LogitsProcessor( + config=SimpleNamespace(model_type="test", vocab_size=8192), + tp_rank=0, + tp_size=2, + tp_group=(0, 1), + ) + lm_head = SimpleNamespace(weight=torch.ones((4096, 2), dtype=torch.bfloat16)) + state = object() + calls = {} + monkeypatch.setattr( + logits_processor_module, + "current_platform", + lambda: SimpleNamespace( + is_nvidia=True, + arch_version=SimpleNamespace(major=major), + ), + ) + monkeypatch.setattr(logits_processor_module, "supports_triton_rsag", lambda: True) + monkeypatch.setattr( + logits_processor_module.pg_manager, + "get_process_group", + lambda backend, group: (backend, group), + ) + + def fake_create_state(**kwargs): + calls.update(kwargs) + return state + + monkeypatch.setattr(logits_processor_module, "create_state", fake_create_state) + + assert processor._init_all_gather_state(lm_head) is state + assert calls == { + "group": ("nccl", (0, 1)), + "rank_in_group": 0, + "max_tokens": processor._LOGITS_AG_MAX_TOKENS, + "hidden_size": 8192, + } + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_fused_softcap_handles_large_logits_without_nan(): cap = 30.0 diff --git a/tokenspeed-kernel/README.md b/tokenspeed-kernel/README.md index 3aaaa5b14..845d9d92f 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -34,7 +34,7 @@ choices (still evolving; subject to change): ### Layered system ``` - public API (mha_prefill, mm, moe_fused, ...) + public API (mha_prefill, mm, mhc_pre, moe_apply, ...) │ ┌───────────┴───────────┐ │ select_kernel │ (family, mode, format_signature, traits, ...) @@ -45,7 +45,7 @@ choices (still evolving; subject to change): └──────────┬──────────┘ │ ┌──────────────┬────────────────┼────────────────┬───────────────┐ - attention gemm moe norm ... (op family) + attention gemm mhc moe ... (op family) │ │ │ │ ┌────┼────┐ ┌────┼────┐ ┌────┼────┐ ┌────┼────┐ triton triton triton triton ← in-tree portable JIT @@ -80,6 +80,7 @@ tokenspeed_kernel/ ops/ attention/ { triton/, flash_attn/, ... } gemm/ { triton.py, trtllm.py, ... } + mhc/ { deep_gemm.py, ... } moe/ { triton.py, deepep.py, triton_kernels.py, ... } ... @@ -132,13 +133,20 @@ backends. See `tokenspeed_kernel/plugins/README.md`. ```python from tokenspeed_kernel import ( mha_prefill, mha_prefill_with_kvcache, mha_decode_with_kvcache, + dsv4_sparse_mla_decode, gdn_chunk_prefill, mm, + mhc_pre, mhc_post, moe_route, moe_dispatch, moe_experts, moe_combine, moe_fused, ... ) ``` +`dsv4_sparse_mla_decode` keeps the model runtime independent of the vendor +library: H100/B200 use the existing FlashMLA solution, while SM120 selects the +FlashInfer 0.6.14 sparse MLA solution over the same packed SWA/compressed-cache +contract. + Using the above platform and solution-agnostic public APIs can get the most value out of TokenSpeed-kernel; but one can also directly call into a specific solution under `ops//`, or manually `select_kernel` with diff --git a/tokenspeed-kernel/python/requirements/cuda.txt b/tokenspeed-kernel/python/requirements/cuda.txt index 1bb05f583..9f5a70600 100644 --- a/tokenspeed-kernel/python/requirements/cuda.txt +++ b/tokenspeed-kernel/python/requirements/cuda.txt @@ -2,7 +2,7 @@ nvidia-cutlass-dsl[cu13]==4.6.0 nvidia-cutlass-dsl-libs-cu13==4.6.0 torch==2.11.0 -flashinfer-python==0.6.13 -flashinfer-cubin==0.6.13 +flashinfer-python==0.6.14 +flashinfer-cubin==0.6.14 nvidia-ml-py==13.580.82 nvtx diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/__init__.py index eb6cedc2f..43889272e 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/__init__.py @@ -31,6 +31,7 @@ dsa_plan, dsa_prefill, dsa_prefill_topk, + dsv4_sparse_mla_decode, gdn_chunk_prefill, mha_decode_with_kvcache, mha_extend_with_kvcache, @@ -38,8 +39,10 @@ mha_prefill, mla_decode_with_kvcache, mla_prefill, + prebuild_dsv4_sparse_mla, ) from tokenspeed_kernel.ops.gemm import mm +from tokenspeed_kernel.ops.mhc import mhc_plan, mhc_post, mhc_pre from tokenspeed_kernel.ops.moe import moe_apply, moe_plan, moe_process_weights from tokenspeed_kernel.ops.quantization import ( quantize_fp8, @@ -57,6 +60,10 @@ "NoKernelFoundError", # gemm "mm", + # mhc + "mhc_plan", + "mhc_pre", + "mhc_post", # attention "mha_plan", "mha_prefill", @@ -66,6 +73,8 @@ "mla_decode_with_kvcache", "dsa_prefill", "dsa_decode", + "dsv4_sparse_mla_decode", + "prebuild_dsv4_sparse_mla", "dsa_prefill_topk", "dsa_decode_topk", "dsa_plan", diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py index 13e38e05e..64e411be6 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py @@ -31,6 +31,9 @@ import tokenspeed_kernel.ops.attention.gluon # noqa: F401 import tokenspeed_kernel.ops.attention.triton # noqa: F401 import torch +from tokenspeed_kernel.ops.attention.flashinfer import ( + prebuild_dsv4_sparse_mla as _prebuild_dsv4_sparse_mla, +) from tokenspeed_kernel.ops.attention.gdn_utils import ( GdnCheckpointLayout, GdnChunkPrefillResult, @@ -65,6 +68,8 @@ def _attention_format_signature(**roles: torch.Tensor): "mla_decode_with_kvcache", "dsa_prefill", "dsa_decode", + "dsv4_sparse_mla_decode", + "prebuild_dsv4_sparse_mla", "dsa_prefill_topk", "dsa_decode_topk", "dsa_plan", @@ -75,6 +80,10 @@ def _attention_format_signature(**roles: torch.Tensor): LSE_LN = math.log2(math.e) +def prebuild_dsv4_sparse_mla() -> None: + _prebuild_dsv4_sparse_mla() + + # ===-----------------------------------------------------------------------===# # GDN Kernels # ===-----------------------------------------------------------------------===# @@ -744,6 +753,131 @@ def mla_decode_with_kvcache( # ===-----------------------------------------------------------------------===# +def dsv4_sparse_mla_decode( + q: torch.Tensor, + swa_kv_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_topk_lens: torch.Tensor, + *, + compressed_kv_cache: torch.Tensor | None = None, + compressed_indices: torch.Tensor | None = None, + compressed_topk_lens: torch.Tensor | None = None, + softmax_scale: float, + sinks: torch.Tensor | None = None, + out: torch.Tensor | None = None, + override: str | None = None, + solution: str | None = None, +) -> torch.Tensor: + """Run DeepSeek V4 sparse MLA over SWA and compressed KV pools. + + FlashInfer 0.6.14 auto-dispatches SM120 calls with more than 64 query + tokens to its sparse MLA prefill kernel; smaller calls use sparse decode. + + Args: + q: BF16 decode or prefill query shaped ``[tokens, heads, 512]``. + swa_kv_cache: Packed ``fp8_ds_mla`` SWA pool shaped + ``[pages, 64, 1, 584]`` (NHD layout). A padded physical page stride + is allowed. + swa_indices: Global SWA slot ids shaped ``[tokens, topk]`` or + ``[tokens, 1, topk]``. + swa_topk_lens: Active SWA entry count per token, shaped ``[tokens]``. + compressed_kv_cache: Optional packed compressed pool in the same NHD + layout. DeepSeek V4 C4A uses 64 rows per page and C128A uses 2. + compressed_indices: Optional global compressed-pool slot ids. + compressed_topk_lens: Active compressed entry count per token. + softmax_scale: Scale applied to QK logits. + sinks: Optional FP32 attention sink, shaped ``[heads]``. + out: Optional BF16 output buffer shaped like ``q``. + override: Optional exact kernel override name. + solution: Optional kernel solution selected through the registry. + + Returns: + BF16 attention output shaped like ``q``. + """ + if q.dim() != 3: + raise ValueError(f"q must be [tokens, heads, 512], got {tuple(q.shape)}") + if swa_kv_cache.dim() != 4 or swa_kv_cache.shape[2] != 1: + raise ValueError( + "swa_kv_cache must use NHD [pages, page_size, 1, bytes] layout, " + f"got {tuple(swa_kv_cache.shape)}" + ) + compressed_args = ( + compressed_kv_cache, + compressed_indices, + compressed_topk_lens, + ) + if any(value is None for value in compressed_args) and any( + value is not None for value in compressed_args + ): + raise ValueError( + "compressed_kv_cache, compressed_indices, and " + "compressed_topk_lens must be provided together" + ) + if compressed_kv_cache is not None and ( + compressed_kv_cache.dim() != 4 or compressed_kv_cache.shape[2] != 1 + ): + raise ValueError( + "compressed_kv_cache must use NHD [pages, page_size, 1, bytes] " + f"layout, got {tuple(compressed_kv_cache.shape)}" + ) + + compressed_page_size = ( + 0 if compressed_kv_cache is None else int(compressed_kv_cache.shape[1]) + ) + traits = { + "head_dim": int(q.shape[-1]), + "swa_page_size": int(swa_kv_cache.shape[1]), + "compressed_page_size": compressed_page_size, + "support_sinks": sinks is not None, + } + signature = _attention_format_signature(q=q) + kernel = select_kernel( + "attention", + "dsv4_sparse_mla_decode", + signature, + traits=traits, + solution=solution, + override=override, + ) + shape_params = { + "tokens": int(q.shape[0]), + "num_heads": int(q.shape[1]), + "head_dim": int(q.shape[2]), + "swa_topk": int(swa_indices.shape[-1]), + "compressed_topk": ( + 0 if compressed_indices is None else int(compressed_indices.shape[-1]) + ), + "swa_page_size": int(swa_kv_cache.shape[1]), + "compressed_page_size": compressed_page_size, + } + ShapeCapture.get().record( + "attention", + "dsv4_sparse_mla_decode", + kernel.name, + q.dtype, + shape_params, + ) + with kernel_scope( + "attention", + "dsv4_sparse_mla_decode", + q.dtype, + kernel_name=kernel.name, + **shape_params, + ): + return kernel( + q=q, + swa_kv_cache=swa_kv_cache, + swa_indices=swa_indices, + swa_topk_lens=swa_topk_lens, + compressed_kv_cache=compressed_kv_cache, + compressed_indices=compressed_indices, + compressed_topk_lens=compressed_topk_lens, + softmax_scale=softmax_scale, + sinks=sinks, + out=out, + ) + + def dsa_decode( q: torch.Tensor, kv_cache: torch.Tensor | None, diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flash_mla/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flash_mla/__init__.py index 74744f623..821413449 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flash_mla/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flash_mla/__init__.py @@ -189,6 +189,7 @@ def _get_decode_sched_meta( solution="flashmla", capability=CapabilityRequirement( min_arch_version=ArchVersion(9, 0), + max_arch_version=ArchVersion(11, 99), vendors=frozenset({"nvidia"}), ), signatures=frozenset({format_signature(q=dense_tensor_format(torch.bfloat16))}), @@ -265,6 +266,89 @@ def flashmla_dsa_decode( return result +if ( + platform.is_nvidia + and platform.is_hopper_plus + and flash_mla_with_kvcache is not error_fn +): + + @register_kernel( + "attention", + "dsv4_sparse_mla_decode", + name="flashmla_dsv4_sparse_mla_decode", + solution="flashmla", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 0), + max_arch_version=ArchVersion(11, 99), + vendors=frozenset({"nvidia"}), + ), + signatures=frozenset({format_signature(q=dense_tensor_format(torch.bfloat16))}), + priority=Priority.PERFORMANT, + traits={ + "head_dim": frozenset({512}), + "swa_page_size": frozenset({64}), + "compressed_page_size": frozenset({0, 2, 64}), + "support_sinks": frozenset({False, True}), + }, + ) + def flashmla_dsv4_sparse_mla_decode( + q: torch.Tensor, + swa_kv_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_topk_lens: torch.Tensor, + compressed_kv_cache: torch.Tensor | None, + compressed_indices: torch.Tensor | None, + compressed_topk_lens: torch.Tensor | None, + softmax_scale: float, + sinks: torch.Tensor | None, + out: torch.Tensor | None, + ) -> torch.Tensor: + if q.dim() != 3: + raise ValueError( + "FlashMLA DeepSeek V4 decode q must be [tokens, heads, dim], " + f"got {tuple(q.shape)}" + ) + q_kernel = q.unsqueeze(1) + swa_indices_kernel = ( + swa_indices.unsqueeze(1) if swa_indices.dim() == 2 else swa_indices + ) + compressed_indices_kernel = ( + compressed_indices.unsqueeze(1) + if compressed_indices is not None and compressed_indices.dim() == 2 + else compressed_indices + ) + result, _ = flash_mla_with_kvcache( + q=q_kernel, + k_cache=swa_kv_cache, + block_table=None, + cache_seqlens=None, + head_dim_v=q.shape[-1], + tile_scheduler_metadata=_get_decode_sched_meta( + q=q_kernel, + num_reqs=q.shape[0], + q_len_per_req=1, + actual_heads=q.shape[1], + topk=swa_indices.shape[-1], + kv_lora_rank=q.shape[-1], + qk_rope_head_dim=64, + ), + softmax_scale=float(softmax_scale), + is_fp8_kvcache=True, + indices=swa_indices_kernel, + attn_sink=sinks, + extra_k_cache=compressed_kv_cache, + extra_indices_in_kvcache=compressed_indices_kernel, + topk_length=swa_topk_lens, + extra_topk_length=compressed_topk_lens, + ) + if result.dim() == 4: + result = result.squeeze(1) + if out is not None: + out.copy_(result) + return out + return result + + if ( platform.is_nvidia and platform.is_hopper_plus diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flashinfer/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flashinfer/__init__.py index bad4db6db..c7885872f 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flashinfer/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flashinfer/__init__.py @@ -46,6 +46,7 @@ trtllm_batch_context_with_kv_cache = error_fn trtllm_batch_decode_with_kv_cache = error_fn trtllm_batch_decode_with_kv_cache_mla = error_fn +trtllm_batch_decode_sparse_mla_dsv4 = error_fn trtllm_ragged_attention_deepseek = error_fn if platform.is_nvidia: @@ -62,13 +63,13 @@ trtllm_ragged_attention_deepseek, ) -if platform.is_blackwell: +if platform.is_blackwell_plus: from flashinfer.mla import ( BatchMLAPagedAttentionWrapper, + trtllm_batch_decode_sparse_mla_dsv4, trtllm_batch_decode_with_kv_cache_mla, ) - # ------------------------------------------------------------------------------ # Kernel registration # ------------------------------------------------------------------------------ @@ -78,6 +79,17 @@ _DSA_SPARSE_WORKSPACE_BYTES = 384 * 1024 * 1024 +def prebuild_dsv4_sparse_mla() -> None: + if not platform.is_nvidia or platform.arch_version.major != 12: + return + if trtllm_batch_decode_sparse_mla_dsv4 is error_fn: + raise RuntimeError("FlashInfer SM120 sparse MLA is unavailable") + + from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module + + get_sparse_mla_sm120_module() + + def _get_dsa_sparse_workspace(device: torch.device | str) -> torch.Tensor: device = torch.device(device) workspace = _dsa_sparse_workspace_buffers.get(device) @@ -125,6 +137,65 @@ def _topk_lens_or_count( return (topk_slots >= 0).sum(dim=-1, dtype=torch.int32).contiguous() +def _flashinfer_dsv4_sparse_mla_decode( + q: torch.Tensor, + swa_kv_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_topk_lens: torch.Tensor, + compressed_kv_cache: torch.Tensor | None, + compressed_indices: torch.Tensor | None, + compressed_topk_lens: torch.Tensor | None, + softmax_scale: float, + sinks: torch.Tensor | None, + out: torch.Tensor | None, +) -> torch.Tensor: + if trtllm_batch_decode_sparse_mla_dsv4 is error_fn: + raise RuntimeError( + "FlashInfer was built without the SM120 DeepSeek V4 sparse MLA API" + ) + return trtllm_batch_decode_sparse_mla_dsv4( + query=q, + swa_kv_cache=swa_kv_cache, + workspace_buffer=_get_dsa_sparse_workspace(q.device), + sparse_indices=swa_indices, + compressed_kv_cache=compressed_kv_cache, + swa_topk_lens=swa_topk_lens, + extra_sparse_indices=compressed_indices, + extra_sparse_topk_lens=compressed_topk_lens, + out=out, + bmm1_scale=float(softmax_scale), + bmm2_scale=1.0, + sinks=sinks, + kv_layout="NHD", + ) + + +if ( + platform.is_nvidia + and platform.arch_version >= ArchVersion(12, 0) + and trtllm_batch_decode_sparse_mla_dsv4 is not error_fn +): + register_kernel( + "attention", + "dsv4_sparse_mla_decode", + name="flashinfer_dsv4_sparse_mla_decode", + solution="flashinfer", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(12, 0), + max_arch_version=ArchVersion(12, 99), + vendors=frozenset({"nvidia"}), + ), + signatures=frozenset({format_signature(q=dense_tensor_format(torch.bfloat16))}), + priority=Priority.SPECIALIZED, + traits={ + "head_dim": frozenset({512}), + "swa_page_size": frozenset({64}), + "compressed_page_size": frozenset({0, 2, 64}), + "support_sinks": frozenset({False, True}), + }, + )(_flashinfer_dsv4_sparse_mla_decode) + + if platform.is_nvidia and platform.is_hopper_plus: @register_kernel( @@ -435,5 +506,6 @@ def flashinfer_trtllm_dsa_prefill( "trtllm_batch_context_with_kv_cache", "trtllm_batch_decode_with_kv_cache", "trtllm_batch_decode_with_kv_cache_mla", + "trtllm_batch_decode_sparse_mla_dsv4", "trtllm_ragged_attention_deepseek", ] diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py new file mode 100644 index 000000000..67fcc208d --- /dev/null +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py @@ -0,0 +1,270 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Manifold-constrained hyper-connection kernel entry points.""" + +from __future__ import annotations + +import math + +import torch +from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope +from tokenspeed_kernel.registry import KernelRegistry +from tokenspeed_kernel.selection import select_kernel +from tokenspeed_kernel.signature import dense_tensor_format, format_signature + +__all__ = ["mhc_plan", "mhc_pre", "mhc_post"] + + +def _require_residual_shape(residual: torch.Tensor) -> tuple[tuple[int, ...], int, int]: + if residual.ndim < 3: + raise ValueError("residual must have shape [..., HC, H]") + outer_shape = tuple(residual.shape[:-2]) + return outer_shape, residual.shape[-2], residual.shape[-1] + + +def mhc_plan( + *, + hc_mult: int, + residual_dtype: torch.dtype = torch.bfloat16, + fn_dtype: torch.dtype = torch.float32, + mix_dtype: torch.dtype = torch.float32, + solution: str | None = None, +) -> dict[str, object]: + """Select matching mHC pre/post kernels for model setup. + + Args: + hc_mult: Number of residual streams. + residual_dtype: Storage dtype of residual and layer tensors. + fn_dtype: Storage dtype of the pre-map projection weight. + mix_dtype: Storage dtype of scale, base, post, and combine tensors. + solution: Optional solution to force through normal selection. + + Returns: + A plan containing the selected pre/post kernel names and their common + solution. Execution can pin either kernel through its recorded name. + """ + pre = select_kernel( + "mhc", + "pre", + format_signature( + residual=dense_tensor_format(residual_dtype), + fn=dense_tensor_format(fn_dtype), + mhc_scale=dense_tensor_format(mix_dtype), + mhc_base=dense_tensor_format(mix_dtype), + ), + traits={"hc_mult": hc_mult}, + solution=solution, + ) + post = select_kernel( + "mhc", + "post", + format_signature( + hidden_states=dense_tensor_format(residual_dtype), + residual=dense_tensor_format(residual_dtype), + post=dense_tensor_format(mix_dtype), + comb=dense_tensor_format(mix_dtype), + ), + traits={"hc_mult": hc_mult}, + solution=solution, + ) + registry = KernelRegistry.get() + pre_spec = registry.get_by_name(pre.name) + post_spec = registry.get_by_name(post.name) + if pre_spec is None or post_spec is None: + raise RuntimeError("Selected mHC kernel is missing its registry spec") + if pre_spec.solution != post_spec.solution: + raise RuntimeError( + "mHC pre/post selection must use one solution, got " + f"{pre_spec.solution!r} and {post_spec.solution!r}" + ) + return { + "hc_mult": hc_mult, + "pre_kernel_name": pre.name, + "post_kernel_name": post.name, + "solution": pre_spec.solution, + } + + +def mhc_pre( + residual: torch.Tensor, + fn: torch.Tensor, + mhc_scale: torch.Tensor, + mhc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + *, + solution: str | None = None, + override: str | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Apply the mHC pre-map through a registered ``mhc.pre`` kernel. + + Args: + residual: Multi-stream residual tensor shaped ``[..., HC, H]``. + fn: Projection weight shaped ``[(2 + HC) * HC, HC * H]``. + mhc_scale: Three FP32 scale values for pre, post, and combine mixes. + mhc_base: FP32 projection bias shaped ``[(2 + HC) * HC]``. + rms_eps: RMS normalization epsilon. + hc_eps: Numerical epsilon used by pre-map and Sinkhorn normalization. + sinkhorn_iters: Number of alternating Sinkhorn normalization iterations. + solution: Optional registered solution to select. + override: Optional exact kernel-name or solution override. + + Returns: + ``(layer_input, post_mix, comb_mix)`` with shapes ``[..., H]``, + ``[..., HC, 1]``, and ``[..., HC, HC]`` respectively. + """ + outer_shape, hc_mult, hidden_size = _require_residual_shape(residual) + num_tokens = math.prod(outer_shape) + if num_tokens == 0: + return ( + residual.new_empty(*outer_shape, hidden_size), + torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ), + torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ), + ) + + signature = format_signature( + residual=dense_tensor_format(residual.dtype), + fn=dense_tensor_format(fn.dtype), + mhc_scale=dense_tensor_format(mhc_scale.dtype), + mhc_base=dense_tensor_format(mhc_base.dtype), + ) + kernel = select_kernel( + "mhc", + "pre", + signature, + traits={"hc_mult": hc_mult}, + solution=solution, + override=override, + ) + shape_params = { + "num_tokens": num_tokens, + "hc_mult": hc_mult, + "hidden_size": hidden_size, + } + ShapeCapture.get().record( + "mhc", + "pre", + kernel.name, + residual.dtype, + shape_params, + ) + with kernel_scope( + "mhc", + "pre", + residual.dtype, + kernel_name=kernel.name, + **shape_params, + ): + return kernel( + residual=residual, + fn=fn, + mhc_scale=mhc_scale, + mhc_base=mhc_base, + rms_eps=rms_eps, + mhc_eps=hc_eps, + sinkhorn_iters=sinkhorn_iters, + ) + + +def mhc_post( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + *, + solution: str | None = None, + override: str | None = None, +) -> torch.Tensor: + """Apply the mHC post-map through a registered ``mhc.post`` kernel. + + Args: + hidden_states: Current layer output shaped ``[..., H]``. + residual: Multi-stream residual tensor shaped ``[..., HC, H]``. + post: Per-stream post mix shaped ``[..., HC, 1]``. + comb: Residual combination matrix shaped ``[..., HC, HC]``. + solution: Optional registered solution to select. + override: Optional exact kernel-name or solution override. + + Returns: + Updated multi-stream residual with the same shape and dtype as + ``residual``. + """ + outer_shape, hc_mult, hidden_size = _require_residual_shape(residual) + num_tokens = math.prod(outer_shape) + if num_tokens == 0: + return torch.empty_like(residual) + + signature = format_signature( + hidden_states=dense_tensor_format(hidden_states.dtype), + residual=dense_tensor_format(residual.dtype), + post=dense_tensor_format(post.dtype), + comb=dense_tensor_format(comb.dtype), + ) + kernel = select_kernel( + "mhc", + "post", + signature, + traits={"hc_mult": hc_mult}, + solution=solution, + override=override, + ) + shape_params = { + "num_tokens": num_tokens, + "hc_mult": hc_mult, + "hidden_size": hidden_size, + } + ShapeCapture.get().record( + "mhc", + "post", + kernel.name, + residual.dtype, + shape_params, + ) + with kernel_scope( + "mhc", + "post", + residual.dtype, + kernel_name=kernel.name, + **shape_params, + ): + return kernel( + hidden_states=hidden_states, + residual=residual, + post=post, + comb=comb, + ) + + +# Backend registration (side-effect imports). +import tokenspeed_kernel.ops.mhc.deep_gemm # noqa: E402,F401 diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/deep_gemm.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/deep_gemm.py new file mode 100644 index 000000000..82592b3e8 --- /dev/null +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/deep_gemm.py @@ -0,0 +1,526 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Portions copyright the vLLM project contributors under Apache-2.0. + +from __future__ import annotations + +from functools import cache + +import torch +from tokenspeed_kernel._triton import tl, triton +from tokenspeed_kernel.platform import ( + ArchVersion, + CapabilityRequirement, + current_platform, +) +from tokenspeed_kernel.registry import Priority, error_fn, register_kernel +from tokenspeed_kernel.signature import dense_tensor_format, format_signature + +deep_gemm = None +if current_platform().is_nvidia: + try: + from tokenspeed_kernel.thirdparty import deep_gemm + except ImportError: + pass + + +def _ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +@cache +def _compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: + device_props = torch.cuda.get_device_properties(0) + split_k = device_props.multi_processor_count // grid_size + if k is not None: + num_block_k = _ceil_div(k, block_k) + split_k = min(split_k, num_block_k // 4) + return max(split_k, 1) + + +@triton.jit +def _load_reduced_mix( + gemm_out_mul, + token_id, + mix_id: tl.constexpr, + num_tokens, + hc_mult3: tl.constexpr, + n_splits: tl.constexpr, +): + value = tl.full((), 0.0, tl.float32) + for split_id in tl.static_range(0, n_splits): + offset = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 + mix_id + value += tl.load(gemm_out_mul + offset) + return value + + +@triton.jit +def _mhc_pre_mix_triton_kernel( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + pre_mix, + post_mix, + comb_mix, + hidden_size: tl.constexpr, + rms_eps: tl.constexpr, + hc_eps: tl.constexpr, + sinkhorn_iters: tl.constexpr, + n_splits: tl.constexpr, + hc_mult: tl.constexpr, + hc_mult2: tl.constexpr, + hc_mult3: tl.constexpr, + block_comb: tl.constexpr, + num_tokens, +): + token_id = tl.program_id(0) + + rms_sum = tl.full((), 0.0, tl.float32) + for split_id in tl.static_range(0, n_splits): + rms_sum += tl.load(gemm_out_sqrsum + split_id * num_tokens + token_id) + rms = tl.rsqrt(rms_sum / (hc_mult * hidden_size) + rms_eps) + + pre_scale = tl.load(hc_scale) + for hc_id in tl.static_range(0, hc_mult): + mix = _load_reduced_mix( + gemm_out_mul, + token_id, + hc_id, + num_tokens, + hc_mult3, + n_splits, + ) + pre = tl.sigmoid(mix * rms * pre_scale + tl.load(hc_base + hc_id)) + hc_eps + tl.store(pre_mix + token_id * hc_mult + hc_id, pre) + + post_scale = tl.load(hc_scale + 1) + for hc_id in tl.static_range(0, hc_mult): + mix = _load_reduced_mix( + gemm_out_mul, + token_id, + hc_mult + hc_id, + num_tokens, + hc_mult3, + n_splits, + ) + post = ( + tl.sigmoid(mix * rms * post_scale + tl.load(hc_base + hc_mult + hc_id)) + * 2.0 + ) + tl.store(post_mix + token_id * hc_mult + hc_id, post) + + comb_offsets = tl.arange(0, block_comb) + comb_mask = comb_offsets < hc_mult2 + comb_scale = tl.load(hc_scale + 2) + comb_mix_values = tl.zeros((block_comb,), tl.float32) + for split_id in tl.static_range(0, n_splits): + split_base = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 + comb_mix_values += tl.load( + gemm_out_mul + split_base + hc_mult * 2 + comb_offsets, + mask=comb_mask, + other=0.0, + ) + comb_values = comb_mix_values * rms * comb_scale + tl.load( + hc_base + hc_mult * 2 + comb_offsets, mask=comb_mask, other=0.0 + ) + rows = comb_offsets // hc_mult + cols = comb_offsets - rows * hc_mult + active = comb_mask + + for row_id in tl.static_range(0, hc_mult): + row_values = tl.where((rows == row_id) & active, comb_values, -float("inf")) + row_max = tl.max(row_values, axis=0) + comb_values = tl.where( + (rows == row_id) & active, tl.exp(comb_values - row_max), comb_values + ) + for row_id in tl.static_range(0, hc_mult): + row_sum = tl.sum(tl.where((rows == row_id) & active, comb_values, 0.0), axis=0) + comb_values = tl.where( + (rows == row_id) & active, comb_values / row_sum + hc_eps, comb_values + ) + for col_id in tl.static_range(0, hc_mult): + col_sum = tl.sum(tl.where((cols == col_id) & active, comb_values, 0.0), axis=0) + comb_values = tl.where( + (cols == col_id) & active, + comb_values / (col_sum + hc_eps), + comb_values, + ) + + for _ in tl.static_range(1, sinkhorn_iters): + for row_id in tl.static_range(0, hc_mult): + row_sum = tl.sum( + tl.where((rows == row_id) & active, comb_values, 0.0), axis=0 + ) + comb_values = tl.where( + (rows == row_id) & active, + comb_values / (row_sum + hc_eps), + comb_values, + ) + for col_id in tl.static_range(0, hc_mult): + col_sum = tl.sum( + tl.where((cols == col_id) & active, comb_values, 0.0), axis=0 + ) + comb_values = tl.where( + (cols == col_id) & active, + comb_values / (col_sum + hc_eps), + comb_values, + ) + + tl.store( + comb_mix + token_id * hc_mult2 + comb_offsets, + comb_values, + mask=comb_mask, + ) + + +@triton.jit +def _mhc_pre_layer_triton_kernel( + pre_mix, + residual, + layer_input, + hidden_size: tl.constexpr, + hc_mult: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + layer_acc = tl.zeros((block_h,), tl.float32) + for hc_id in tl.static_range(0, hc_mult): + pre = tl.load(pre_mix + token_id * hc_mult + hc_id).to(tl.float32) + residual_offsets = ( + token_id * hc_mult * hidden_size + hc_id * hidden_size + hidden_offsets + ) + residual_values = tl.load( + residual + residual_offsets, mask=hidden_mask, other=0.0 + ).to(tl.float32) + layer_acc += pre * residual_values + tl.store( + layer_input + token_id * hidden_size + hidden_offsets, + layer_acc, + mask=hidden_mask, + ) + + +@triton.jit +def _mhc_post_triton_kernel( + comb, + residual, + post, + hidden_states, + out, + hidden_size: tl.constexpr, + hc_mult: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + hidden_values = tl.load( + hidden_states + token_id * hidden_size + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + for out_hc in tl.static_range(0, hc_mult): + acc = tl.load(post + token_id * hc_mult + out_hc).to(tl.float32) * hidden_values + for in_hc in tl.static_range(0, hc_mult): + comb_value = tl.load( + comb + token_id * hc_mult * hc_mult + in_hc * hc_mult + out_hc + ).to(tl.float32) + residual_values = tl.load( + residual + + token_id * hc_mult * hidden_size + + in_hc * hidden_size + + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + acc += comb_value * residual_values + tl.store( + out + + token_id * hc_mult * hidden_size + + out_hc * hidden_size + + hidden_offsets, + acc, + mask=hidden_mask, + ) + + +@triton.jit +def _mhc_post_hc4_triton_kernel( + comb, + residual, + post, + hidden_states, + out, + hidden_size: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + token_hidden_offset = token_id * hidden_size + token_residual_offset = token_id * 4 * hidden_size + + hidden_values = tl.load( + hidden_states + token_hidden_offset + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + post_base = token_id * 4 + acc0 = tl.load(post + post_base + 0).to(tl.float32) * hidden_values + acc1 = tl.load(post + post_base + 1).to(tl.float32) * hidden_values + acc2 = tl.load(post + post_base + 2).to(tl.float32) * hidden_values + acc3 = tl.load(post + post_base + 3).to(tl.float32) * hidden_values + + comb_base = token_id * 16 + for in_hc in tl.static_range(0, 4): + residual_values = tl.load( + residual + token_residual_offset + in_hc * hidden_size + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + comb_row = comb_base + in_hc * 4 + acc0 += tl.load(comb + comb_row + 0).to(tl.float32) * residual_values + acc1 += tl.load(comb + comb_row + 1).to(tl.float32) * residual_values + acc2 += tl.load(comb + comb_row + 2).to(tl.float32) * residual_values + acc3 += tl.load(comb + comb_row + 3).to(tl.float32) * residual_values + + tl.store( + out + token_residual_offset + hidden_offsets, + acc0, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size + hidden_offsets, + acc1, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size * 2 + hidden_offsets, + acc2, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size * 3 + hidden_offsets, + acc3, + mask=hidden_mask, + ) + + +def _deep_gemm_mhc_pre( + residual: torch.Tensor, + fn: torch.Tensor, + mhc_scale: torch.Tensor, + mhc_base: torch.Tensor, + rms_eps: float, + mhc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if residual.dtype != torch.bfloat16 or fn.dtype != torch.float32: + raise RuntimeError("fast mHC requires bf16 residual and fp32 weights") + if not residual.is_cuda: + raise RuntimeError("fast mHC requires CUDA tensors") + + if deep_gemm is None: + raise RuntimeError("deep_gemm.tf32_hc_prenorm_gemm is unavailable") + + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + hc_mult2 = hc_mult * hc_mult + hc_mult3 = hc_mult * 2 + hc_mult2 + hc_hidden_size = hc_mult * hidden_size + outer_shape = residual.shape[:-2] + residual_flat = residual.view(-1, hc_mult, hidden_size) + num_tokens = residual_flat.shape[0] + if num_tokens == 0: + return ( + residual.new_empty(*outer_shape, hidden_size), + torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ), + torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ), + ) + + block_k = 64 + block_m = 64 + n_splits = _compute_num_split( + block_k, hc_hidden_size, _ceil_div(num_tokens, block_m) + ) + + post_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + pre_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + comb_mix = torch.empty( + num_tokens, hc_mult2, dtype=torch.float32, device=residual.device + ) + layer_input = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ) + gemm_out_mul = torch.empty( + n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device + ) + gemm_out_sqrsum = torch.empty( + n_splits, num_tokens, dtype=torch.float32, device=residual.device + ) + + deep_gemm.tf32_hc_prenorm_gemm( + residual_flat.view(num_tokens, hc_hidden_size), + fn, + gemm_out_mul, + gemm_out_sqrsum, + n_splits, + ) + block_h = 1024 + block_comb = triton.next_power_of_2(hc_mult2) + _mhc_pre_mix_triton_kernel[(num_tokens,)]( + gemm_out_mul, + gemm_out_sqrsum, + mhc_scale, + mhc_base, + pre_mix, + post_mix, + comb_mix, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_eps=mhc_eps, + sinkhorn_iters=sinkhorn_iters, + n_splits=n_splits, + hc_mult=hc_mult, + hc_mult2=hc_mult2, + hc_mult3=hc_mult3, + block_comb=block_comb, + num_tokens=num_tokens, + num_warps=1, + ) + _mhc_pre_layer_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + pre_mix, + residual_flat, + layer_input, + hidden_size=hidden_size, + hc_mult=hc_mult, + block_h=block_h, + num_warps=4, + ) + + return ( + layer_input.view(*outer_shape, hidden_size), + post_mix.view(*outer_shape, hc_mult, 1), + comb_mix.view(*outer_shape, hc_mult, hc_mult), + ) + + +def _deep_gemm_mhc_post( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + if not hidden_states.is_cuda: + raise RuntimeError("fast mHC requires CUDA tensors") + if residual.numel() == 0: + return torch.empty_like(residual) + out = torch.empty_like(residual) + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + residual_flat = residual.view(-1, hc_mult, hidden_size) + hidden_states_flat = hidden_states.view(-1, hidden_size) + post_flat = post.view(-1, hc_mult) + comb_flat = comb.view(-1, hc_mult, hc_mult) + num_tokens = residual_flat.shape[0] + if hc_mult == 4: + block_h = 256 + _mhc_post_hc4_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + comb_flat, + residual_flat, + post_flat, + hidden_states_flat, + out, + hidden_size=hidden_size, + block_h=block_h, + num_warps=4, + ) + return out + + block_h = 1024 + _mhc_post_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + comb_flat, + residual_flat, + post_flat, + hidden_states_flat, + out, + hidden_size=hidden_size, + hc_mult=hc_mult, + block_h=block_h, + num_warps=4, + ) + return out + + +deep_gemm_mhc_pre = error_fn +deep_gemm_mhc_post = error_fn + +if deep_gemm is not None: + capability = CapabilityRequirement( + vendors=frozenset({"nvidia"}), + min_arch_version=ArchVersion(9, 0), + max_arch_version=ArchVersion(12, 99), + ) + deep_gemm_mhc_pre = register_kernel( + "mhc", + "pre", + name="deep_gemm_mhc_pre", + solution="deep_gemm", + capability=capability, + signatures={ + format_signature( + residual=dense_tensor_format(torch.bfloat16), + fn=dense_tensor_format(torch.float32), + mhc_scale=dense_tensor_format(torch.float32), + mhc_base=dense_tensor_format(torch.float32), + ) + }, + traits={"hc_mult": frozenset({4})}, + priority=Priority.SPECIALIZED, + tags={"latency"}, + )(_deep_gemm_mhc_pre) + deep_gemm_mhc_post = register_kernel( + "mhc", + "post", + name="deep_gemm_mhc_post", + solution="deep_gemm", + capability=capability, + signatures={ + format_signature( + hidden_states=dense_tensor_format(torch.bfloat16), + residual=dense_tensor_format(torch.bfloat16), + post=dense_tensor_format(torch.float32), + comb=dense_tensor_format(torch.float32), + ) + }, + traits={"hc_mult": frozenset({4})}, + priority=Priority.SPECIALIZED, + tags={"latency"}, + )(_deep_gemm_mhc_post) + +__all__ = ["deep_gemm_mhc_pre", "deep_gemm_mhc_post"] diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/__init__.py index f65a7ff6a..b2756eac3 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/__init__.py @@ -20,6 +20,7 @@ import tokenspeed_kernel.ops.moe.flashinfer.cutedsl_deepep_nvfp4 # noqa: F401 import tokenspeed_kernel.ops.moe.flashinfer.cutlass_fp8 # noqa: F401 +import tokenspeed_kernel.ops.moe.flashinfer.cutlass_mxfp4 # noqa: F401 import tokenspeed_kernel.ops.moe.flashinfer.cutlass_nvfp4 # noqa: F401 import tokenspeed_kernel.ops.moe.flashinfer.cutlass_unquant # noqa: F401 import tokenspeed_kernel.ops.moe.flashinfer.trtllm_fp8 # noqa: F401 diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/cutlass_mxfp4.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/cutlass_mxfp4.py new file mode 100644 index 000000000..8887662b2 --- /dev/null +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/flashinfer/cutlass_mxfp4.py @@ -0,0 +1,283 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.platform import ( + ArchVersion, + CapabilityRequirement, + current_platform, +) +from tokenspeed_kernel.registry import Priority, register_kernel +from tokenspeed_kernel.signature import format_signatures + +platform = current_platform() +next_power_of_2 = lambda value: 1 if value <= 1 else 1 << (value - 1).bit_length() + + +def _reorder_w13(x: torch.Tensor, layout: str, dim: int) -> torch.Tensor: + if dim < 0: + dim += x.dim() + size = x.shape[dim] + if size % 2 != 0: + raise ValueError(f"expected even size in dim {dim}, got {size}") + if layout == "concatenated": + first, second = x.split(size // 2, dim=dim) + return torch.cat((second, first), dim=dim).contiguous() + if layout == "interleaved": + shape = list(x.shape) + paired_shape = shape[:dim] + [size // 2, 2] + shape[dim + 1 :] + paired = x.reshape(paired_shape) + w1 = paired.select(dim + 1, 0) + w3 = paired.select(dim + 1, 1) + return torch.cat((w3, w1), dim=dim).contiguous() + raise ValueError(f"unknown w13_input_layout: {layout!r}") + + +def _swizzle_mxfp4_block_scales(scales: torch.Tensor) -> torch.Tensor: + """Convert linear E8M0 scales to CUTLASS' padded 128x4 SF layout.""" + if scales.ndim != 3 or scales.dtype != torch.uint8: + raise ValueError( + "MXFP4 scales must be a 3D uint8 tensor with shape " + f"[experts, rows, cols], got {tuple(scales.shape)} {scales.dtype}" + ) + + experts, rows, cols = scales.shape + rows_padded = (rows + 127) // 128 * 128 + cols_padded = (cols + 3) // 4 * 4 + padded = torch.zeros( + (experts, rows_padded, cols_padded), + dtype=scales.dtype, + device=scales.device, + ) + padded[:, :rows, :cols] = scales + return ( + padded.reshape( + experts, + rows_padded // 128, + 4, + 32, + cols_padded // 4, + 4, + ) + .permute(0, 1, 4, 3, 2, 5) + .contiguous() + .reshape(experts, rows_padded, cols_padded) + ) + + +def _expert_float_parameter( + w: torch.nn.Module, + value: float | None, +) -> torch.nn.Parameter | None: + if value is None: + return None + return torch.nn.Parameter( + torch.full( + (w.w13_weight.shape[0],), + float(value), + dtype=torch.float32, + device=w.w13_weight.device, + ), + requires_grad=False, + ) + + +if platform.is_nvidia: + from flashinfer import ActivationType, cutlass_fused_moe, mxfp8_quantize + from flashinfer.fused_moe.core import get_cutlass_fused_moe_module + from flashinfer.quantization.fp8_quantization import ( + get_mxfp8_quantization_sm100_module, + ) + + def flashinfer_cutlass_mxfp4_moe_weights( + plan: dict, + w: torch.nn.Module, + ) -> None: + """Prepare checkpoint MXFP4 weights for FlashInfer CUTLASS on SM12x.""" + del plan + intermediate_size = w.w13_weight.shape[1] // 2 + hidden_size = w.w2_weight.shape[1] + if intermediate_size % 128 != 0 or hidden_size % 128 != 0: + raise ValueError( + "FlashInfer SM12x MXFP4 MoE requires hidden and intermediate " + f"sizes divisible by 128, got {hidden_size} and {intermediate_size}" + ) + + # FlashInfer's CUTLASS SwiGLU path consumes w3 before w1. + w13_layout = getattr(w, "w13_input_layout", "concatenated") + w13_weight = _reorder_w13(w.w13_weight.data, w13_layout, 1) + w13_scale = _reorder_w13(w.w13_weight_scale.data, w13_layout, 1) + w2_weight = w.w2_weight.data.contiguous() + w2_scale = w.w2_weight_scale.data.contiguous() + + w.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + w.w13_weight_scale = torch.nn.Parameter( + _swizzle_mxfp4_block_scales(w13_scale), + requires_grad=False, + ) + w.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + w.w2_weight_scale = torch.nn.Parameter( + _swizzle_mxfp4_block_scales(w2_scale), + requires_grad=False, + ) + + if hasattr(w, "w13_weight_bias"): + w.w13_weight_bias = torch.nn.Parameter( + _reorder_w13(w.w13_weight_bias.data, w13_layout, 1), + requires_grad=False, + ) + if hasattr(w, "w2_weight_bias"): + w.w2_weight_bias = torch.nn.Parameter( + w.w2_weight_bias.data.contiguous(), + requires_grad=False, + ) + + num_local_experts = w.w13_weight.shape[0] + ones = torch.ones( + num_local_experts, + dtype=torch.float32, + device=w.w13_weight.device, + ) + w.w13_weight_global_scale = torch.nn.Parameter(ones, requires_grad=False) + w.w2_weight_global_scale = torch.nn.Parameter(ones.clone(), requires_grad=False) + + swiglu_arg = getattr(w, "swiglu_arg", None) + w.gemm1_alpha = _expert_float_parameter( + w, None if swiglu_arg is None else swiglu_arg.alpha + ) + w.gemm1_beta = _expert_float_parameter(w, getattr(w, "swiglu_beta", None)) + w.gemm1_clamp_limit = _expert_float_parameter( + w, None if swiglu_arg is None else swiglu_arg.limit + ) + w.intermediate_size_per_partition = intermediate_size + w.hidden_size_padded = hidden_size + w.hidden_size_original = getattr(w, "hidden_size", hidden_size) + + if w.w13_weight.is_cuda: + major, minor = torch.cuda.get_device_capability(w.w13_weight.device) + get_cutlass_fused_moe_module(f"{major}{minor}") + get_mxfp8_quantization_sm100_module() + + @register_kernel( + "moe", + "apply", + name="flashinfer_cutlass_mxfp4_moe_apply", + solution="flashinfer_cutlass", + weight_preprocessor=flashinfer_cutlass_mxfp4_moe_weights, + capability=CapabilityRequirement( + vendors=frozenset({"nvidia"}), + min_arch_version=ArchVersion(12, 0), + max_arch_version=ArchVersion(12, 1), + ), + signatures=format_signatures( + "x", + "dense", + {torch.float16, torch.bfloat16}, + ), + traits={ + "weight_dtype": frozenset({"mxfp4"}), + "activation": frozenset({"swiglu"}), + "routing_mode": frozenset({"precomputed_topk"}), + "supports_deferred_finalize": frozenset({False}), + "supports_ep": frozenset({True}), + "supports_all_to_all_ep": frozenset({False}), + "ispp_alignment": frozenset({128}), + "internal_activation_dtype": frozenset({"input"}), + "supports_bias": frozenset({True}), + }, + priority=Priority.PERFORMANT, + ) + def flashinfer_cutlass_mxfp4_moe_apply( + plan: dict, + x: torch.Tensor, + w: torch.nn.Module, + router_logits: torch.Tensor, + topk_weights: torch.Tensor | None = None, + topk_ids: torch.Tensor | None = None, + num_tokens_global: int | None = None, + max_num_tokens_per_gpu: int | None = None, + do_finalize: bool = True, + enable_pdl: bool = False, + ) -> torch.Tensor: + del plan, router_logits, num_tokens_global, max_num_tokens_per_gpu + if not do_finalize: + raise ValueError("FlashInfer CUTLASS MXFP4 MoE requires finalization") + if topk_weights is None or topk_ids is None: + raise ValueError( + "FlashInfer CUTLASS MXFP4 MoE requires precomputed top-k routing" + ) + + hidden_padded = getattr(w, "hidden_size_padded", w.w2_weight.shape[1]) + hidden_original = getattr(w, "hidden_size_original", hidden_padded) + if x.shape[0] == 0: + return x.new_empty((0, hidden_original), dtype=torch.bfloat16) + if x.shape[-1] > hidden_padded: + raise ValueError( + f"input hidden size {x.shape[-1]} exceeds prepared size {hidden_padded}" + ) + if x.shape[-1] < hidden_padded: + x = torch.nn.functional.pad(x, (0, hidden_padded - x.shape[-1])) + + x_quant, x_scale = mxfp8_quantize( + x, + False, + alignment=hidden_padded, + enable_pdl=enable_pdl, + ) + output = torch.empty( + (x.shape[0], hidden_padded), + dtype=torch.bfloat16, + device=x.device, + ) + result = cutlass_fused_moe( + input=x_quant, + input_sf=x_scale, + swizzled_input_sf=False, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights.to(torch.float32), + fc1_expert_weights=w.w13_weight.view(torch.long), + fc2_expert_weights=w.w2_weight.view(torch.long), + fc1_expert_biases=getattr(w, "w13_weight_bias", None), + fc2_expert_biases=getattr(w, "w2_weight_bias", None), + output_dtype=torch.bfloat16, + output=output, + quant_scales=[ + w.w13_weight_scale.contiguous().view(torch.int32), + w.w13_weight_global_scale, + w.w2_weight_scale.contiguous().view(torch.int32), + w.w2_weight_global_scale, + ], + swiglu_alpha=getattr(w, "gemm1_alpha", None), + swiglu_beta=getattr(w, "gemm1_beta", None), + swiglu_limit=getattr(w, "gemm1_clamp_limit", None), + ep_size=getattr(w, "ep_size", 1), + ep_rank=getattr(w, "ep_rank", 0), + tp_size=getattr(w, "tp_size", 1), + tp_rank=getattr(w, "tp_rank", 0), + use_mxfp8_act_scaling=True, + tune_max_num_tokens=next_power_of_2(x.shape[0]), + enable_pdl=enable_pdl, + activation_type=ActivationType.Swiglu, + )[0] + if hidden_original != hidden_padded: + result = result[:, :hidden_original].contiguous() + return result diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/registry.py b/tokenspeed-kernel/python/tokenspeed_kernel/registry.py index b44969836..1f4c91c7b 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/registry.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/registry.py @@ -482,6 +482,7 @@ def load_builtin_kernels() -> None: del sys.modules[key] import tokenspeed_kernel.ops.embedding # noqa: F401 import tokenspeed_kernel.ops.gemm # noqa: F401 + import tokenspeed_kernel.ops.mhc # noqa: F401 import tokenspeed_kernel.ops.moe # noqa: F401 import tokenspeed_kernel.ops.quantization # noqa: F401 import tokenspeed_kernel.ops.sampling # noqa: F401 diff --git a/tokenspeed-kernel/test/conftest.py b/tokenspeed-kernel/test/conftest.py index 97c3c19c9..943536ea9 100644 --- a/tokenspeed-kernel/test/conftest.py +++ b/tokenspeed-kernel/test/conftest.py @@ -197,6 +197,33 @@ def b200_platform() -> PlatformInfo: ) +@pytest.fixture +def sm120_platform() -> PlatformInfo: + return PlatformInfo( + vendor="nvidia", + arch_version=ArchVersion(12, 0), + device_name="NVIDIA RTX PRO 6000 Blackwell Server Edition", + device_count=4, + total_memory=96 * (1024**3), + memory_bandwidth=1792.0, + sm_count=188, + max_threads_per_sm=1536, + max_shared_memory_per_sm=102400, + sm_features=frozenset( + { + "tensor_core:f16", + "tensor_core:int8", + "tensor_core:f8", + "tensor_core:f4", + "memory:async_copy", + "memory:tma", + } + ), + runtime_features=frozenset({"runtime:cuda_graph"}), + interconnect=InterconnectInfo(topology="pcie"), + ) + + @pytest.fixture def fresh_registry(): KernelRegistry.reset() diff --git a/tokenspeed-kernel/test/ops/test_attention_dsv4_sparse_mla.py b/tokenspeed-kernel/test/ops/test_attention_dsv4_sparse_mla.py new file mode 100644 index 000000000..f3aa03345 --- /dev/null +++ b/tokenspeed-kernel/test/ops/test_attention_dsv4_sparse_mla.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +import importlib + +import pytest +import tokenspeed_kernel.ops.attention as attention_ops +import tokenspeed_kernel.ops.attention.flash_mla as flashmla_ops +import tokenspeed_kernel.ops.attention.flashinfer as flashinfer_ops +import torch +from tokenspeed_kernel.platform import Platform +from tokenspeed_kernel.registry import KernelRegistry, error_fn +from tokenspeed_kernel.selection import select_kernel +from tokenspeed_kernel.signature import dense_tensor_format, format_signature + + +class _SelectedKernel: + name = "test_dsv4_sparse_mla" + + def __init__(self, result: torch.Tensor) -> None: + self.result = result + self.calls: list[dict[str, object]] = [] + + def __call__(self, **kwargs: object) -> torch.Tensor: + self.calls.append(kwargs) + return self.result + + +def _inputs() -> dict[str, torch.Tensor]: + return { + "q": torch.randn(2, 64, 512, dtype=torch.bfloat16), + "swa_kv_cache": torch.empty(3, 64, 1, 584, dtype=torch.uint8), + "swa_indices": torch.zeros(2, 128, dtype=torch.int32), + "swa_topk_lens": torch.full((2,), 128, dtype=torch.int32), + "compressed_kv_cache": torch.empty(4, 2, 1, 584, dtype=torch.uint8), + "compressed_indices": torch.zeros(2, 1, 512, dtype=torch.int32), + "compressed_topk_lens": torch.full((2,), 512, dtype=torch.int32), + "sinks": torch.zeros(64, dtype=torch.float32), + } + + +def test_dsv4_sparse_mla_decode_dispatches_through_kernel_boundary( + monkeypatch, +) -> None: + inputs = _inputs() + expected = torch.empty_like(inputs["q"]) + selected = _SelectedKernel(expected) + selection: dict[str, object] = {} + + def fake_select_kernel(family, mode, signature, **kwargs): + selection.update( + family=family, + mode=mode, + signature=signature, + kwargs=kwargs, + ) + return selected + + monkeypatch.setattr(attention_ops, "select_kernel", fake_select_kernel) + + actual = attention_ops.dsv4_sparse_mla_decode( + inputs["q"], + inputs["swa_kv_cache"], + inputs["swa_indices"], + inputs["swa_topk_lens"], + compressed_kv_cache=inputs["compressed_kv_cache"], + compressed_indices=inputs["compressed_indices"], + compressed_topk_lens=inputs["compressed_topk_lens"], + softmax_scale=512**-0.5, + sinks=inputs["sinks"], + solution="flashinfer", + ) + + assert actual is expected + assert selection["family"] == "attention" + assert selection["mode"] == "dsv4_sparse_mla_decode" + signature = selection["signature"] + assert signature.storage_dtype_for("q") == torch.bfloat16 + assert selection["kwargs"] == { + "traits": { + "head_dim": 512, + "swa_page_size": 64, + "compressed_page_size": 2, + "support_sinks": True, + }, + "solution": "flashinfer", + "override": None, + } + assert selected.calls == [ + { + "q": inputs["q"], + "swa_kv_cache": inputs["swa_kv_cache"], + "swa_indices": inputs["swa_indices"], + "swa_topk_lens": inputs["swa_topk_lens"], + "compressed_kv_cache": inputs["compressed_kv_cache"], + "compressed_indices": inputs["compressed_indices"], + "compressed_topk_lens": inputs["compressed_topk_lens"], + "softmax_scale": 512**-0.5, + "sinks": inputs["sinks"], + "out": None, + } + ] + + +def test_flashinfer_dsv4_sparse_mla_uses_dual_cache_sm120_api(monkeypatch) -> None: + inputs = _inputs() + expected = torch.empty_like(inputs["q"]) + workspace = torch.empty(1, dtype=torch.uint8) + call: dict[str, object] = {} + + def fake_sparse_mla(**kwargs): + call.update(kwargs) + return expected + + monkeypatch.setattr( + flashinfer_ops, + "trtllm_batch_decode_sparse_mla_dsv4", + fake_sparse_mla, + ) + monkeypatch.setattr( + flashinfer_ops, + "_get_dsa_sparse_workspace", + lambda _device: workspace, + ) + + actual = flashinfer_ops._flashinfer_dsv4_sparse_mla_decode( + q=inputs["q"], + swa_kv_cache=inputs["swa_kv_cache"], + swa_indices=inputs["swa_indices"], + swa_topk_lens=inputs["swa_topk_lens"], + compressed_kv_cache=inputs["compressed_kv_cache"], + compressed_indices=inputs["compressed_indices"], + compressed_topk_lens=inputs["compressed_topk_lens"], + softmax_scale=512**-0.5, + sinks=inputs["sinks"], + out=None, + ) + + assert actual is expected + assert call == { + "query": inputs["q"], + "swa_kv_cache": inputs["swa_kv_cache"], + "workspace_buffer": workspace, + "sparse_indices": inputs["swa_indices"], + "compressed_kv_cache": inputs["compressed_kv_cache"], + "swa_topk_lens": inputs["swa_topk_lens"], + "extra_sparse_indices": inputs["compressed_indices"], + "extra_sparse_topk_lens": inputs["compressed_topk_lens"], + "out": None, + "bmm1_scale": 512**-0.5, + "bmm2_scale": 1.0, + "sinks": inputs["sinks"], + "kv_layout": "NHD", + } + + +@pytest.mark.parametrize( + "platform_fixture,expected_solution", + [ + ("h100_platform", "flashmla"), + ("b200_platform", "flashmla"), + ("sm120_platform", "flashinfer"), + ], +) +def test_dsv4_sparse_mla_solution_is_architecture_scoped( + platform_fixture: str, + expected_solution: str, + request, +) -> None: + platform = request.getfixturevalue(platform_fixture) + real_platform = Platform.get() + try: + Platform.override(platform) + KernelRegistry.reset() + importlib.reload(flashmla_ops) + importlib.reload(flashinfer_ops) + if ( + expected_solution == "flashmla" + and flashmla_ops.flash_mla_with_kvcache is error_fn + ): + pytest.skip("FlashMLA is unavailable") + if ( + expected_solution == "flashinfer" + and flashinfer_ops.trtllm_batch_decode_sparse_mla_dsv4 is error_fn + ): + pytest.skip("FlashInfer SM120 sparse MLA is unavailable") + + kernel = select_kernel( + "attention", + "dsv4_sparse_mla_decode", + format_signature(q=dense_tensor_format(torch.bfloat16)), + traits={ + "head_dim": 512, + "swa_page_size": 64, + "compressed_page_size": 2, + "support_sinks": True, + }, + ) + + assert kernel.name == f"{expected_solution}_dsv4_sparse_mla_decode" + finally: + Platform.override(real_platform) + KernelRegistry.reset() + importlib.reload(flashmla_ops) + importlib.reload(flashinfer_ops) diff --git a/tokenspeed-kernel/test/ops/test_mhc.py b/tokenspeed-kernel/test/ops/test_mhc.py new file mode 100644 index 000000000..8335c9c93 --- /dev/null +++ b/tokenspeed-kernel/test/ops/test_mhc.py @@ -0,0 +1,316 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +import importlib + +import pytest +import tokenspeed_kernel.ops.mhc as mhc_ops +import tokenspeed_kernel.ops.mhc.deep_gemm as deep_gemm_mhc +import torch +import torch.nn.functional as F +from tokenspeed_kernel.platform import Platform, current_platform +from tokenspeed_kernel.registry import KernelRegistry + + +class _SelectedKernel: + name = "test_mhc_kernel" + + def __init__(self, result: object) -> None: + self.result = result + self.calls: list[dict[str, object]] = [] + + def __call__(self, **kwargs: object) -> object: + self.calls.append(kwargs) + return self.result + + +@pytest.mark.parametrize( + "platform_fixture", + [ + "h100_platform", + "b200_platform", + "sm120_platform", + ], +) +def test_mhc_plan_preserves_existing_architecture_solutions( + platform_fixture: str, + request, + fresh_registry, +) -> None: + platform = request.getfixturevalue(platform_fixture) + real_platform = Platform.get() + try: + Platform.override(platform) + importlib.reload(deep_gemm_mhc) + KernelRegistry.get().clear_cache() + if deep_gemm_mhc.deep_gemm is None: + pytest.skip("DeepGEMM is unavailable") + + plan = mhc_ops.mhc_plan(hc_mult=4) + + assert plan == { + "hc_mult": 4, + "pre_kernel_name": "deep_gemm_mhc_pre", + "post_kernel_name": "deep_gemm_mhc_post", + "solution": "deep_gemm", + } + finally: + Platform.override(real_platform) + KernelRegistry.get().clear_cache() + + +def test_mhc_pre_dispatches_through_registered_kernel(monkeypatch) -> None: + residual = torch.randn(2, 4, 8, dtype=torch.bfloat16) + fn = torch.randn(24, 32, dtype=torch.float32) + scale = torch.ones(3, dtype=torch.float32) + base = torch.zeros(24, dtype=torch.float32) + expected = ( + torch.empty(2, 8, dtype=torch.bfloat16), + torch.empty(2, 4, 1, dtype=torch.float32), + torch.empty(2, 4, 4, dtype=torch.float32), + ) + selected = _SelectedKernel(expected) + selection: dict[str, object] = {} + + def fake_select_kernel(family, mode, signature, **kwargs): + selection.update( + family=family, + mode=mode, + signature=signature, + kwargs=kwargs, + ) + return selected + + monkeypatch.setattr(mhc_ops, "select_kernel", fake_select_kernel) + + actual = mhc_ops.mhc_pre( + residual, + fn, + scale, + base, + rms_eps=1e-6, + hc_eps=2e-6, + sinkhorn_iters=3, + solution="deep_gemm", + ) + + assert actual is expected + assert selection["family"] == "mhc" + assert selection["mode"] == "pre" + signature = selection["signature"] + assert signature.storage_dtype_for("residual") == torch.bfloat16 + assert signature.storage_dtype_for("fn") == torch.float32 + assert signature.storage_dtype_for("mhc_scale") == torch.float32 + assert signature.storage_dtype_for("mhc_base") == torch.float32 + assert selection["kwargs"] == { + "traits": {"hc_mult": 4}, + "solution": "deep_gemm", + "override": None, + } + assert selected.calls == [ + { + "residual": residual, + "fn": fn, + "mhc_scale": scale, + "mhc_base": base, + "rms_eps": 1e-6, + "mhc_eps": 2e-6, + "sinkhorn_iters": 3, + } + ] + + +def test_mhc_post_dispatches_through_registered_kernel(monkeypatch) -> None: + hidden_states = torch.randn(2, 8, dtype=torch.bfloat16) + residual = torch.randn(2, 4, 8, dtype=torch.bfloat16) + post = torch.randn(2, 4, 1, dtype=torch.float32) + comb = torch.randn(2, 4, 4, dtype=torch.float32) + expected = torch.empty_like(residual) + selected = _SelectedKernel(expected) + selection: dict[str, object] = {} + + def fake_select_kernel(family, mode, signature, **kwargs): + selection.update( + family=family, + mode=mode, + signature=signature, + kwargs=kwargs, + ) + return selected + + monkeypatch.setattr(mhc_ops, "select_kernel", fake_select_kernel) + + actual = mhc_ops.mhc_post( + hidden_states, + residual, + post, + comb, + solution="deep_gemm", + ) + + assert actual is expected + assert selection["family"] == "mhc" + assert selection["mode"] == "post" + signature = selection["signature"] + assert signature.storage_dtype_for("hidden_states") == torch.bfloat16 + assert signature.storage_dtype_for("residual") == torch.bfloat16 + assert signature.storage_dtype_for("post") == torch.float32 + assert signature.storage_dtype_for("comb") == torch.float32 + assert selection["kwargs"] == { + "traits": {"hc_mult": 4}, + "solution": "deep_gemm", + "override": None, + } + assert selected.calls == [ + { + "hidden_states": hidden_states, + "residual": residual, + "post": post, + "comb": comb, + } + ] + + +def test_mhc_pre_empty_input_preserves_output_contract(monkeypatch) -> None: + residual = torch.empty(0, 4, 8, dtype=torch.bfloat16) + fn = torch.empty(24, 32, dtype=torch.float32) + scale = torch.ones(3, dtype=torch.float32) + base = torch.zeros(24, dtype=torch.float32) + + def fail_select(*args, **kwargs): + raise AssertionError("empty input must not dispatch a kernel") + + monkeypatch.setattr(mhc_ops, "select_kernel", fail_select) + + layer_input, post, comb = mhc_ops.mhc_pre( + residual, + fn, + scale, + base, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=2, + ) + + assert layer_input.shape == (0, 8) + assert layer_input.dtype == torch.bfloat16 + assert post.shape == (0, 4, 1) + assert post.dtype == torch.float32 + assert comb.shape == (0, 4, 4) + assert comb.dtype == torch.float32 + + +def _mhc_reference( + residual: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + *, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + tokens, hc_mult, hidden_size = residual.shape + x = residual.reshape(tokens, hc_mult * hidden_size).float() + rms = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x, fn) * rms + pre_raw, post_raw, comb_raw = torch.split( + mixes, [hc_mult, hc_mult, hc_mult * hc_mult], dim=-1 + ) + pre_base, post_base, comb_base = torch.split( + base, [hc_mult, hc_mult, hc_mult * hc_mult] + ) + pre = torch.sigmoid(pre_raw * scale[0] + pre_base) + hc_eps + post = (torch.sigmoid(post_raw * scale[1] + post_base) * 2.0).unsqueeze(-1) + comb = torch.softmax( + comb_raw.reshape(tokens, hc_mult, hc_mult) * scale[2] + + comb_base.reshape(1, hc_mult, hc_mult), + dim=-1, + ) + comb = comb + hc_eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps) + layer_input = torch.sum(pre.unsqueeze(-1) * residual.float(), dim=1) + return layer_input.to(torch.bfloat16), post, comb + + +@pytest.mark.skipif( + not current_platform().is_nvidia, + reason="mHC GPU kernels require an NVIDIA GPU", +) +def test_mhc_pre_and_post_match_reference( + device: str, + fresh_registry, +) -> None: + importlib.reload(deep_gemm_mhc) + if deep_gemm_mhc.deep_gemm is None: + pytest.skip("DeepGEMM mHC is unavailable") + + from tokenspeed_kernel import mhc_post, mhc_pre + + torch.manual_seed(1234) + residual = torch.randn(3, 4, 64, device=device, dtype=torch.bfloat16) + fn = torch.randn(24, 256, device=device, dtype=torch.float32) + scale = torch.tensor([0.7, 1.1, 0.5], device=device, dtype=torch.float32) + base = torch.randn(24, device=device, dtype=torch.float32) + + actual = mhc_pre( + residual, + fn, + scale, + base, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=3, + solution="deep_gemm", + ) + expected = _mhc_reference( + residual, + fn, + scale, + base, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=3, + ) + + torch.testing.assert_close(actual[0], expected[0], rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual[1], expected[1], rtol=1e-2, atol=5e-3) + torch.testing.assert_close(actual[2], expected[2], rtol=1e-2, atol=5e-3) + + updated = mhc_post( + actual[0], + residual, + actual[1], + actual[2], + solution="deep_gemm", + ) + expected_updated = torch.einsum("tnm,tnh->tmh", actual[2].float(), residual.float()) + expected_updated += actual[1].float() * actual[0].float().unsqueeze(1) + torch.testing.assert_close( + updated, + expected_updated.to(torch.bfloat16), + rtol=2e-2, + atol=2e-2, + ) diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index b223fccb3..dbe1fb258 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/tokenspeed-kernel/test/test_kernel_api_selection.py @@ -24,6 +24,7 @@ import importlib from dataclasses import dataclass +from types import SimpleNamespace from typing import Callable import pytest @@ -82,6 +83,7 @@ cutedsl_deepep_nvfp4 as _moe_cutedsl_deepep_nvfp4, ) from tokenspeed_kernel.ops.moe.flashinfer import cutlass_fp8 as _moe_cutlass_fp8 +from tokenspeed_kernel.ops.moe.flashinfer import cutlass_mxfp4 as _moe_cutlass_mxfp4 from tokenspeed_kernel.ops.moe.flashinfer import cutlass_nvfp4 as _moe_cutlass_nvfp4 from tokenspeed_kernel.ops.moe.flashinfer import cutlass_unquant as _moe_cutlass_unquant from tokenspeed_kernel.ops.moe.flashinfer import trtllm_fp8 as _moe_trtllm_fp8 @@ -125,6 +127,7 @@ # MoE registration modules. _moe_cutedsl_deepep_nvfp4, _moe_cutlass_fp8, + _moe_cutlass_mxfp4, _moe_cutlass_nvfp4, _moe_cutlass_unquant, _moe_trtllm_fp8, @@ -200,6 +203,125 @@ def preprocess(plan, w): assert calls == [(plan, module)] +def test_mxfp4_cutlass_plan_is_available_on_sm120( + sm120_platform: PlatformInfo, +) -> None: + if not hasattr(_moe_cutlass_mxfp4, "flashinfer_cutlass_mxfp4_moe_apply"): + pytest.skip("FlashInfer CUTLASS MXFP4 is NVIDIA-only") + + registry = KernelRegistry.get() + real_platform = Platform.get() + try: + Platform.override(sm120_platform) + registry.clear_cache() + + plan = tokenspeed_kernel.moe_plan( + "mxfp4", + input_dtype=torch.bfloat16, + activation="swiglu", + ep_size=4, + ispp=18432, + internal_activation_dtype="input", + with_bias=True, + solution="flashinfer_cutlass", + ) + + _assert_moe_plan( + plan, + apply="flashinfer_cutlass_mxfp4_moe_apply", + preprocessor="flashinfer_cutlass_mxfp4_moe_weights", + ) + finally: + Platform.override(real_platform) + registry.clear_cache() + + +@pytest.mark.parametrize("w13_layout", ["concatenated", "interleaved"]) +def test_mxfp4_cutlass_preprocessor_preserves_checkpoint_values( + w13_layout: str, +) -> None: + if not hasattr(_moe_cutlass_mxfp4, "flashinfer_cutlass_mxfp4_moe_weights"): + pytest.skip("FlashInfer CUTLASS MXFP4 is NVIDIA-only") + + module = torch.nn.Module() + module.hidden_size = 128 + module.w13_weight = torch.nn.Parameter( + torch.arange(1 * 256 * 64, dtype=torch.int64) + .remainder(256) + .to(torch.uint8) + .reshape(1, 256, 64), + requires_grad=False, + ) + module.w13_weight_scale = torch.nn.Parameter( + torch.arange(1 * 256 * 4, dtype=torch.int64) + .remainder(256) + .to(torch.uint8) + .reshape(1, 256, 4), + requires_grad=False, + ) + module.w2_weight = torch.nn.Parameter( + torch.zeros((1, 128, 64), dtype=torch.uint8), requires_grad=False + ) + module.w2_weight_scale = torch.nn.Parameter( + torch.arange(1 * 128 * 4, dtype=torch.int64) + .remainder(256) + .to(torch.uint8) + .reshape(1, 128, 4), + requires_grad=False, + ) + module.w13_weight_bias = torch.nn.Parameter( + torch.arange(256, dtype=torch.bfloat16).reshape(1, 256), + requires_grad=False, + ) + module.w2_weight_bias = torch.nn.Parameter( + torch.zeros((1, 128), dtype=torch.bfloat16), requires_grad=False + ) + module.swiglu_arg = SimpleNamespace(alpha=None, limit=7.0) + module.swiglu_beta = None + module.w13_input_layout = w13_layout + + original_weight = module.w13_weight.detach().clone() + original_scale = module.w13_weight_scale.detach().clone() + original_bias = module.w13_weight_bias.detach().clone() + _moe_cutlass_mxfp4.flashinfer_cutlass_mxfp4_moe_weights({}, module) + + if w13_layout == "concatenated": + expected_weight = torch.cat( + (original_weight[:, 128:], original_weight[:, :128]), dim=1 + ) + expected_scale = torch.cat( + (original_scale[:, 128:], original_scale[:, :128]), dim=1 + ) + expected_bias = torch.cat( + (original_bias[:, 128:], original_bias[:, :128]), dim=1 + ) + else: + expected_weight = torch.cat( + (original_weight[:, 1::2], original_weight[:, 0::2]), dim=1 + ) + expected_scale = torch.cat( + (original_scale[:, 1::2], original_scale[:, 0::2]), dim=1 + ) + expected_bias = torch.cat( + (original_bias[:, 1::2], original_bias[:, 0::2]), dim=1 + ) + + assert torch.equal(module.w13_weight, expected_weight) + assert torch.equal(module.w13_weight_bias, expected_bias) + unswizzled_scale = ( + module.w13_weight_scale.reshape(1, 2, 1, 32, 4, 4) + .permute(0, 1, 4, 3, 2, 5) + .reshape(1, 256, 4) + ) + assert torch.equal(unswizzled_scale, expected_scale) + assert module.w13_weight_scale.view(torch.int32).shape == (1, 256, 1) + assert torch.equal(module.w13_weight_global_scale, torch.ones(1)) + assert torch.equal(module.w2_weight_global_scale, torch.ones(1)) + assert module.gemm1_alpha is None + assert module.gemm1_beta is None + assert torch.equal(module.gemm1_clamp_limit, torch.tensor([7.0])) + + @dataclass(frozen=True) class KernelApiSelectionCase: id: str diff --git a/tokenspeed-kernel/test/test_runtime_callsite_selection.py b/tokenspeed-kernel/test/test_runtime_callsite_selection.py index 22d2bf764..5e696642a 100644 --- a/tokenspeed-kernel/test/test_runtime_callsite_selection.py +++ b/tokenspeed-kernel/test/test_runtime_callsite_selection.py @@ -46,6 +46,7 @@ cutedsl_deepep_nvfp4 as _moe_cutedsl_deepep_nvfp4, ) from tokenspeed_kernel.ops.moe.flashinfer import cutlass_fp8 as _moe_cutlass_fp8 +from tokenspeed_kernel.ops.moe.flashinfer import cutlass_mxfp4 as _moe_cutlass_mxfp4 from tokenspeed_kernel.ops.moe.flashinfer import cutlass_nvfp4 as _moe_cutlass_nvfp4 from tokenspeed_kernel.ops.moe.flashinfer import cutlass_unquant as _moe_cutlass_unquant from tokenspeed_kernel.ops.moe.flashinfer import trtllm_mxfp4 as _moe_trtllm_mxfp4 @@ -68,6 +69,7 @@ # MoE _moe_cutedsl_deepep_nvfp4, _moe_cutlass_fp8, + _moe_cutlass_mxfp4, _moe_cutlass_nvfp4, _moe_cutlass_unquant, _moe_trtllm_mxfp4, diff --git a/tokenspeed-kernel/test/thirdparty/test_deep_gemm_warmup.py b/tokenspeed-kernel/test/thirdparty/test_deep_gemm_warmup.py new file mode 100644 index 000000000..c656e3413 --- /dev/null +++ b/tokenspeed-kernel/test/thirdparty/test_deep_gemm_warmup.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.thirdparty.deep_gemm import warmup + + +def test_prefill_warmup_keeps_deep_gemm_mhc_enabled_by_default(monkeypatch) -> None: + calls: list[tuple[list[dict], int, torch.device]] = [] + + def fake_prenorm_warmup(shapes, max_tokens, device): + calls.append((shapes, max_tokens, device)) + + monkeypatch.setattr( + warmup, + "_warmup_tf32_hc_prenorm_gemm", + fake_prenorm_warmup, + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + device = torch.device("cuda", 0) + + warmup.warmup_prefill_jit( + hidden_size=64, + num_attention_heads=8, + hc_mult=4, + max_tokens=32, + device=device, + ) + + assert calls == [ + ( + [{"hc_hidden_size": 256, "mix_hc": 24, "hc_dim": 256}], + 32, + device, + ) + ]