feat(deepseek-v4): enable SM120 serving

.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 }} \

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)
 

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