feat: SM120 (Blackwell Desktop) support for DeepSeek-V4 inference (#24692)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx

@@ -33,6 +33,7 @@ export const DeepSeekV4Deployment = () => {
         { id: "gb300", label: "GB300", default: false },
         { id: "h200",  label: "H200",  default: false },
         { id: "h100",  label: "H100",  default: false },
+        { id: "sm120", label: "RTX PRO 6000 (SM120)", default: false },
       ],
     },
     modelSize: {
@@ -124,7 +125,7 @@ export const DeepSeekV4Deployment = () => {
   // low-latency / balanced / cp recipes, and on PD-Disagg (the cookbook's
   // PD command builder doesn't emit the megamoe backend / env vars yet).
   const MEGAMOE_UNSUPPORTED_RECIPES = new Set(["low-latency", "balanced", "cp", "pd-disagg"]);
-  const MEGAMOE_UNSUPPORTED_HARDWARE = new Set(["h100", "h200"]);
+  const MEGAMOE_UNSUPPORTED_HARDWARE = new Set(["h100", "h200", "sm120"]);
   const isMegamoeUnsupported = (vals) =>
     MEGAMOE_UNSUPPORTED_HARDWARE.has(vals.hardware) ||
     MEGAMOE_UNSUPPORTED_RECIPES.has(vals.recipe);
@@ -184,6 +185,23 @@ export const DeepSeekV4Deployment = () => {
           : it
       );
     }
+    // SM120: Flash only (Pro doesn't fit in 8× 96 GB).
+    if (option.name === "modelSize" && vals && vals.hardware === "sm120") {
+      return option.items.map((it) =>
+        it.id === "big"
+          ? { ...it, disabled: true, disabledReason: "V4-Pro does not fit on SM120 (8× 96 GB)" }
+          : it
+      );
+    }
+    // SM120: TP-only, no EP / CP / PD-Disagg.
+    if (option.name === "recipe" && vals && vals.hardware === "sm120") {
+      const sm120Unsupported = new Set(["balanced", "max-throughput", "cp", "pd-disagg"]);
+      return option.items.map((it) =>
+        sm120Unsupported.has(it.id)
+          ? { ...it, disabled: true, disabledReason: "SM120 supports low-latency (TP-only) recipe" }
+          : it
+      );
+    }
     if (option.name === "modelSize" && vals && isProDisabledFp8H100(vals)) {
       return option.items.map((it) =>
         it.id === "big"
@@ -342,6 +360,9 @@ export const DeepSeekV4Deployment = () => {
     // the generator. TP=8 single-node uses the same sgl-project FP8 ckpt as
     // H200; the Flash/balanced/max-throughput recipes use TP=8 DP=8 + DeepEP.
     "h100-fp8|small": { slug: "sgl-project/DeepSeek-V4-Flash-FP8", tp: 8, multinode: false },
+    // SM120 (RTX PRO 6000): Flash only, TP=4 single-node. Uses Marlin MoE runner
+    // with SM120 Triton fallback kernels. Requires lmsysorg/sglang:dev-cu13.
+    "sm120|small": { slug: "deepseek-ai/DeepSeek-V4-Flash", tp: 4, multinode: false },
   };
   // Per (hardware, modelSize) PD role TP (from allinone _PD_SPEC).
   const PD_TP_SPEC = {
@@ -412,6 +433,7 @@ export const DeepSeekV4Deployment = () => {
     "h100-fp8|small|low-latency",
     "h100-fp8|small|balanced",
     "h100-fp8|small|max-throughput",
+    "sm120|small|low-latency",
   ]);
   // Recipes whose command is intentionally not yet provided (e.g. blocked by an
   // upstream limitation). Showing a minimal placeholder is friendlier to users
@@ -475,6 +497,30 @@ export const DeepSeekV4Deployment = () => {
       return buildPDDisaggCommand(hardware, modelSize);
     }
 
+    // SM120 (RTX PRO 6000) path: Flash only, TP=4, Marlin MoE w/ SM120 Triton
+    // fallback. Requires Docker image lmsysorg/sglang:dev-cu13 (CUDA 13).
+    if (hardware === "sm120") {
+      const verifyKey = `${hardware}|${modelSize}|${recipe}`;
+      const sm120Flags = [
+        "  --trust-remote-code",
+        `  --model-path ${slug}`,
+        `  --tp ${tp}`,
+        "  --moe-runner-backend marlin",
+        "  --mem-fraction-static 0.70",
+        "  --cuda-graph-max-bs 32",
+      ];
+      if (toolcall === "enabled") sm120Flags.push("  --tool-call-parser deepseekv4");
+      if (reasoningParser === "enabled") sm120Flags.push("  --reasoning-parser deepseek-v4");
+      sm120Flags.push("  --host 0.0.0.0");
+      sm120Flags.push("  --port 30000");
+
+      const sm120Note = "# SM120: use Docker image lmsysorg/sglang:dev-cu13\n";
+      const sm120Cmd = `${sm120Note}sglang serve \\\n${sm120Flags.join(" \\\n")}`;
+      return VERIFIED_RECIPES.has(verifyKey)
+        ? sm120Cmd
+        : `${BEING_VERIFIED_NOTE}\n${commentOutCommand(sm120Cmd)}`;
+    }
+
     // H200 (FP4) path: dedicated branch — Hopper runs the FP4-mixed Instruct
     // repos through one of two w4a16 MoE runners (Marlin or Flashinfer mxfp4),
     // so it doesn't share envs or flags with either the FP8 H200 path or the

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -56,6 +56,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import ceil_align
+from sglang.srt.utils.common import is_sm120_supported
 
 if TYPE_CHECKING:
     from sgl_kernel.flash_mla import FlashMLASchedMeta
@@ -63,6 +64,8 @@ if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
 
+_is_sm120 = is_sm120_supported()
+
 logger = logging.getLogger(__name__)
 
 SWA_WINDOW = 128
@@ -82,6 +85,8 @@ def _pad_last_dim(x: T, multiples_of: int = PAGE_INDEX_ALIGNED_SIZE) -> T:
 
 
 def _create_flashmla_metadata():
+    if _is_sm120:
+        return None
     import sgl_kernel.flash_mla as flash_mla
 
     return flash_mla.get_mla_metadata()[0]
@@ -1042,24 +1047,42 @@ class DeepseekV4AttnBackend(
                     extra_indices.shape[-1] % 64 == 0
                 ), f"{extra_indices.shape=}'s last dimension is not aligned to 64"
 
-            import sgl_kernel.flash_mla as flash_mla
+            if _is_sm120:
+                from sglang.srt.layers.attention.flash_mla_sm120 import (
+                    flash_mla_with_kvcache_sm120,
+                )
 
-            o = flash_mla.flash_mla_with_kvcache(
+                o = flash_mla_with_kvcache_sm120(
-                q=q,
+                    q=q,
-                k_cache=swa_k_cache,
+                    k_cache=swa_k_cache,
-                head_dim_v=self.head_dim_v,
+                    head_dim_v=self.head_dim_v,
-                block_table=None,
+                    softmax_scale=self.softmax_scale,
-                cache_seqlens=None,
+                    indices=swa_page_indices,
-                tile_scheduler_metadata=flashmla_metadata,
+                    topk_length=swa_topk_lengths,
-                softmax_scale=self.softmax_scale,
+                    attn_sink=attn_sink,
-                is_fp8_kvcache=True,
+                    extra_k_cache=extra_k_cache,
-                indices=swa_page_indices,
+                    extra_indices_in_kvcache=extra_indices,
-                topk_length=swa_topk_lengths,
+                    extra_topk_length=extra_topk_lengths,
-                attn_sink=attn_sink,
+                )[0]
-                extra_k_cache=extra_k_cache,
+            else:
-                extra_indices_in_kvcache=extra_indices,
+                import sgl_kernel.flash_mla as flash_mla
-                extra_topk_length=extra_topk_lengths,
+
-            )[0]
+                o = flash_mla.flash_mla_with_kvcache(
+                    q=q,
+                    k_cache=swa_k_cache,
+                    head_dim_v=self.head_dim_v,
+                    block_table=None,
+                    cache_seqlens=None,
+                    tile_scheduler_metadata=flashmla_metadata,
+                    softmax_scale=self.softmax_scale,
+                    is_fp8_kvcache=True,
+                    indices=swa_page_indices,
+                    topk_length=swa_topk_lengths,
+                    attn_sink=attn_sink,
+                    extra_k_cache=extra_k_cache,
+                    extra_indices_in_kvcache=extra_indices,
+                    extra_topk_length=extra_topk_lengths,
+                )[0]
 
             o = o.squeeze(1)
             return o

python/sglang/srt/layers/attention/dsv4/indexer.py

@@ -20,6 +20,7 @@ from sglang.srt.layers.attention.dsv4.metadata import PagedIndexerMetadata
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
 from sglang.srt.utils import add_prefix, is_hip
+from sglang.srt.utils.common import is_sm120_supported
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -148,6 +149,74 @@ def _aiter_fp8_paged_mqa_logits(
     return logits
 
 
+def fp8_paged_mqa_logits_torch_sm120(
+    q_fp8: torch.Tensor,
+    kvcache_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    deep_gemm_metadata: Any,
+    max_seq_len: int,
+    clean_logits: bool = True,
+) -> torch.Tensor:
+    """CUDA-graph-compatible FP8 paged MQA logits for SM120 (vectorized, no .item())."""
+    _ = deep_gemm_metadata
+    batch_size, _, num_heads, head_dim = q_fp8.shape
+    block_size = kvcache_fp8.shape[1]
+    device = q_fp8.device
+
+    assert head_dim == 128, "Vectorized torch impl hardcodes DSV4 indexer head_dim=128"
+    assert (
+        block_size == 64
+    ), "Vectorized torch impl hardcodes block_size=64 cache layout"
+    assert q_fp8.shape == (batch_size, 1, num_heads, head_dim)
+    assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4)
+    assert weight.shape == (batch_size, num_heads)
+    if seq_lens.dim() > 1:
+        seq_lens = seq_lens.squeeze(-1)
+    assert seq_lens.shape == (batch_size,)
+    assert page_table.shape[0] == batch_size
+    assert clean_logits == False
+
+    max_pages = (max_seq_len + block_size - 1) // block_size
+    max_padded_seq = max_pages * block_size
+
+    kvcache_flat = kvcache_fp8.view(-1, block_size * (head_dim + 4))
+    SCALE_OFFSET = block_size * head_dim
+
+    page_ids = page_table[:, :max_pages]
+    kvcache_gathered = kvcache_flat[page_ids]
+
+    kv_value_raw = kvcache_gathered[..., :SCALE_OFFSET]
+    kv_scale_raw = kvcache_gathered[..., SCALE_OFFSET:]
+
+    kv_value = kv_value_raw.contiguous().view(dtype=FP8_DTYPE).to(torch.float32)
+    kv_value = kv_value.view(batch_size, max_padded_seq, head_dim)
+
+    kv_scale = kv_scale_raw.contiguous().view(dtype=torch.float32)
+    kv_scale = kv_scale.view(batch_size, max_padded_seq)
+
+    q = q_fp8[:, 0].to(torch.float32)
+
+    score = torch.bmm(kv_value, q.transpose(1, 2))
+
+    score = F.relu(score)
+    score = score * weight.unsqueeze(1)
+    score = score.sum(dim=2)
+
+    score = score * kv_scale
+
+    out_width = min(max_padded_seq, max_seq_len)
+    logits = score.new_full((batch_size, max_seq_len), float("-inf"))
+    logits[:, :out_width] = score[:, :out_width]
+
+    positions = torch.arange(max_seq_len, device=device)
+    invalid_mask = positions.unsqueeze(0) >= seq_lens.unsqueeze(1)
+    logits.masked_fill_(invalid_mask, float("-inf"))
+
+    return logits
+
+
 def topk_transform_512_pytorch_vectorized(
     scores: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -428,7 +497,10 @@ class C4IndexerBackendMixin:
         elif envs.SGLANG_OPT_USE_AITER_INDEXER.get():
             fn = _aiter_fp8_paged_mqa_logits
         elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
-            fn = fp8_paged_mqa_logits_torch
+            if is_sm120_supported():
+                fn = fp8_paged_mqa_logits_torch_sm120
+            else:
+                fn = fp8_paged_mqa_logits_torch
         else:
             from deep_gemm import fp8_paged_mqa_logits as fn
 

python/sglang/srt/layers/attention/flash_mla_sm120.py

@@ -0,0 +1,252 @@
+"""SM120 FlashMLA sparse decode implementation.
+
+On SM120 (Blackwell Desktop / RTX PRO 6000) the flash_mla CUDA kernel
+is not available, so this module provides alternative implementations:
+
+- A fused Triton kernel (default, ``SGLANG_SM120_TRITON_FLASHMLA=1``)
+- A pure-PyTorch fallback (``SGLANG_SM120_TRITON_FLASHMLA=0``)
+
+The FP8 KV cache uses a page-internal layout where NOPE+ROPE data has
+stride (nope_dim + rope_dim*2) per token, and scales are stored in a
+separate region at the end of each page.
+"""
+
+import logging
+import os
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Page layout constants for DSv4-Flash (MODEL1):
+#   nope_dim = 448, rope_dim = 64, quantize_block_size = 64
+#   nope_rope_stride = 448 + 64*2 = 576 bytes per token
+#   scale_stride = ceil(448/64) + 1 = 8 bytes per token (7 scales + 1 pad)
+#   bytes_per_token = 448 + 128 + 8 = 584
+#   page_bytes = ceil_div(page_size * 584, 576) * 576
+
+_NOPE_DIM = 448
+_ROPE_DIM = 64
+_NOPE_ROPE_STRIDE = _NOPE_DIM + _ROPE_DIM * 2  # 576
+_TILE_SIZE = 64
+_NUM_TILES = _NOPE_DIM // _TILE_SIZE  # 7
+_SCALE_STRIDE = _NUM_TILES + 1  # 8 (7 scales + 1 pad)
+_D = _NOPE_DIM + _ROPE_DIM  # 512
+
+
+def _gather_and_dequant(k_cache, indices, page_size):
+    """Gather KV entries from the paged buffer using correct page-internal addressing.
+
+    Args:
+        k_cache: (num_pages, page_size, 1, bytes_per_token) float8_e4m3fn
+                 Non-contiguous view of the raw page buffer.
+        indices: (...) int32/int64, token-level indices. -1 = invalid.
+        page_size: tokens per page (256)
+
+    Returns:
+        kv: (..., _D) bfloat16, dequantized KV vectors
+    """
+    idx_shape = indices.shape
+    flat_idx = indices.reshape(-1)  # (N,)
+    N = flat_idx.shape[0]
+    device = k_cache.device
+
+    # Page-level addressing
+    page_bytes = k_cache.stride(0)  # actual byte stride between pages
+    pages = flat_idx // page_size
+    offsets = flat_idx % page_size
+
+    # Clamp invalid indices
+    safe_pages = pages.clamp(min=0)
+    safe_offsets = offsets.clamp(min=0)
+
+    # Access raw buffer as uint8 — use as_strided to get full page view
+    num_pages = k_cache.shape[0]
+    raw_pages = k_cache.as_strided(
+        (num_pages, page_bytes),
+        (page_bytes, 1),
+    ).view(
+        torch.uint8
+    )  # (num_pages, page_bytes) uint8
+    # Note: float8_e4m3fn and uint8 are both 1 byte, view is safe
+
+    # Compute byte offsets within each page
+    # NOPE: page[safe_page, safe_offset * 576 + 0:448]
+    # ROPE: page[safe_page, safe_offset * 576 + 448:576]
+    # SCALES: page[safe_page, page_size * 576 + safe_offset * 8 + 0:7]
+
+    nope_base = safe_offsets * _NOPE_ROPE_STRIDE  # (N,)
+    nope_offsets = nope_base.unsqueeze(-1) + torch.arange(
+        _NOPE_DIM, device=device, dtype=torch.long
+    )  # (N, 448)
+
+    rope_base = nope_base + _NOPE_DIM  # (N,)
+    rope_offsets = rope_base.unsqueeze(-1) + torch.arange(
+        _ROPE_DIM * 2, device=device, dtype=torch.long
+    )  # (N, 128)
+
+    scale_section_offset = page_size * _NOPE_ROPE_STRIDE  # 147456
+    scale_base = scale_section_offset + safe_offsets * _SCALE_STRIDE  # (N,)
+    scale_offsets = scale_base.unsqueeze(-1) + torch.arange(
+        _NUM_TILES, device=device, dtype=torch.long
+    )  # (N, 7)
+
+    # Gather bytes per page — use advanced indexing
+    # raw_pages[safe_pages, nope_offsets] → (N, 448)
+    page_idx_nope = safe_pages.unsqueeze(-1).expand_as(nope_offsets)
+    nope_bytes = raw_pages[page_idx_nope, nope_offsets]  # (N, 448) uint8
+
+    page_idx_rope = safe_pages.unsqueeze(-1).expand_as(rope_offsets)
+    rope_bytes = raw_pages[page_idx_rope, rope_offsets]  # (N, 128) uint8
+
+    page_idx_scale = safe_pages.unsqueeze(-1).expand_as(scale_offsets)
+    scale_bytes = raw_pages[page_idx_scale, scale_offsets]  # (N, 7) uint8
+
+    # Reinterpret dtypes
+    nope_fp8 = nope_bytes.view(torch.float8_e4m3fn)  # (N, 448)
+    rope_bf16 = rope_bytes.contiguous().view(torch.bfloat16)  # (N, 64)
+    scale_e8m0 = scale_bytes.view(torch.float8_e8m0fnu)  # (N, 7)
+
+    # Dequantize: nope_tile * scale_tile → bf16 (vectorized)
+    result = torch.empty(N, _D, dtype=torch.bfloat16, device=device)
+    result[:, :_NOPE_DIM] = (
+        (
+            nope_fp8.view(N, _NUM_TILES, _TILE_SIZE).float()
+            * scale_e8m0.view(N, _NUM_TILES, 1).float()
+        )
+        .view(N, _NOPE_DIM)
+        .to(torch.bfloat16)
+    )
+    result[:, _NOPE_DIM:] = rope_bf16
+
+    return result.reshape(*idx_shape, _D)
+
+
+def _sm120_sparse_decode_fwd(
+    q,
+    k_cache,
+    indices,
+    topk_length,
+    attn_sink,
+    head_dim_v,
+    softmax_scale,
+    extra_k_cache=None,
+    extra_indices=None,
+    extra_topk_length=None,
+):
+    B, s_q, H_q, D_qk = q.shape
+    num_pages, page_size, H_k, bpt = k_cache.shape
+    topk = indices.shape[-1]
+
+    invalid_mask = indices < 0
+    safe_indices = indices.clamp(min=0)
+
+    if topk_length is not None:
+        topk_range = torch.arange(topk, device=topk_length.device).view(1, 1, topk)
+        invalid_mask = invalid_mask | (topk_range >= topk_length.view(B, 1, 1))
+
+    # Gather and dequantize using page-aware addressing
+    gathered_kv = _gather_and_dequant(k_cache, safe_indices, page_size)
+
+    if extra_k_cache is not None and extra_indices is not None:
+        extra_topk = extra_indices.shape[-1]
+        extra_page_size = extra_k_cache.shape[1]
+        extra_invalid = extra_indices < 0
+        extra_safe = extra_indices.clamp(min=0)
+        if extra_topk_length is not None:
+            extra_range = torch.arange(
+                extra_topk, device=extra_topk_length.device
+            ).view(1, 1, extra_topk)
+            extra_invalid = extra_invalid | (
+                extra_range >= extra_topk_length.view(B, 1, 1)
+            )
+        extra_kv = _gather_and_dequant(extra_k_cache, extra_safe, extra_page_size)
+        gathered_kv = torch.cat([gathered_kv, extra_kv], dim=2)
+        invalid_mask = torch.cat([invalid_mask, extra_invalid], dim=2)
+
+    gathered_kv[invalid_mask] = 0.0
+
+    q_f = q.float()
+    kv_f = gathered_kv.float()
+    kv_d = kv_f.shape[-1]
+    if D_qk != kv_d:
+        q_f = q_f[..., :kv_d]
+
+    scores = torch.einsum("bshd,bstd->bsht", q_f, kv_f) * softmax_scale
+    scores.masked_fill_(invalid_mask.unsqueeze(2).expand_as(scores), float("-inf"))
+
+    lse = torch.logsumexp(scores, dim=-1)
+
+    if attn_sink is not None:
+        lse_for_out = torch.logsumexp(
+            torch.stack([lse, attn_sink.view(1, 1, H_q).expand_as(lse)], dim=0), dim=0
+        )
+    else:
+        lse_for_out = lse.clone()
+
+    lonely = lse == float("-inf")
+    lse_for_out[lonely] = float("inf")
+    weights = torch.exp(scores - lse_for_out.unsqueeze(-1))
+    out = torch.einsum("bsht,bstv->bshv", weights, kv_f[..., :head_dim_v])
+    out[lonely.unsqueeze(-1).expand_as(out)] = 0.0
+
+    return out.to(torch.bfloat16), lse.permute(0, 2, 1)
+
+
+# Default SM120 FlashMLA backend: "triton" (optimized) or "torch" (pure-PyTorch fallback).
+# Controlled by SGLANG_SM120_TRITON_FLASHMLA env var (1=triton, 0=torch).
+_sm120_default_backend = (
+    "triton" if os.environ.get("SGLANG_SM120_TRITON_FLASHMLA", "1") == "1" else "torch"
+)
+
+
+def flash_mla_with_kvcache_sm120(**kwargs):
+    """SM120 FlashMLA sparse decode entry point.
+
+    Dispatches to the Triton kernel (default) or PyTorch fallback.
+    """
+    q = kwargs["q"]
+    k_cache = kwargs["k_cache"]
+    indices = kwargs["indices"]
+    topk_length = kwargs.get("topk_length")
+    attn_sink = kwargs.get("attn_sink")
+    head_dim_v = kwargs["head_dim_v"]
+    softmax_scale = kwargs.get("softmax_scale")
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    extra_k_cache = kwargs.get("extra_k_cache")
+    extra_indices = kwargs.get("extra_indices_in_kvcache")
+    extra_topk_length = kwargs.get("extra_topk_length")
+
+    if _sm120_default_backend == "triton":
+        from sglang.srt.layers.attention.flash_mla_sm120_triton import (
+            flash_mla_sparse_decode_triton,
+        )
+
+        out, lse = flash_mla_sparse_decode_triton(
+            q,
+            k_cache,
+            indices,
+            topk_length,
+            attn_sink,
+            head_dim_v,
+            softmax_scale,
+            extra_k_cache,
+            extra_indices,
+            extra_topk_length,
+        )
+        return (out, lse)
+
+    out, lse = _sm120_sparse_decode_fwd(
+        q,
+        k_cache,
+        indices,
+        topk_length,
+        attn_sink,
+        head_dim_v,
+        softmax_scale,
+        extra_k_cache,
+        extra_indices,
+        extra_topk_length,
+    )
+    return (out, lse)

python/sglang/srt/layers/attention/flash_mla_sm120_triton.py

@@ -0,0 +1,370 @@
+"""SM120-optimized Triton FlashMLA sparse decode kernel — Tiled V2.
+
+Replaces V1's serial token loop with a tiled vectorized approach:
+  1. BLOCK_T tokens loaded simultaneously via 2D gather (vs 1-at-a-time)
+  2. All BLOCK_T QK scores computed at once via vectorized mul-reduce
+  3. V accumulation via vectorized weighted sum across BLOCK_T tokens
+  4. Online softmax operates on tile-level maxima (fewer rescales)
+
+Three typed views of the same paged buffer handle FP8/uint8/BF16 regions:
+- float8_e4m3fn view -> nope FP8 values (direct load + dequant)
+- uint8 view -> UE8M0 scale bytes (raw integer -> exp2 conversion)
+- bfloat16 view -> rope BF16 values (direct load)
+
+DSv4 page layout (per token, 576 bytes data + 8 bytes scales):
+  Data section: [0:448] FP8 nope | [448:576] BF16 rope (64 values = 128 bytes)
+  Scale section: [page_size*576 + offset*8 : +7] UE8M0 scales (7 groups of 64)
+
+Target: RTX PRO 6000 (SM120, 188 SMs, 99KB SMEM, ~1.5 TB/s GDDR7, 96MB L2)
+"""
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+LOG2E = tl.constexpr(1.4426950408889634)
+
+# DSv4 KV cache layout constants
+_NOPE_DIM = 448
+_ROPE_DIM = 64
+_D = _NOPE_DIM + _ROPE_DIM  # 512
+_TOKEN_DATA_STRIDE = 576  # bytes per token in data section
+_SCALE_STRIDE = 8  # bytes per token in scale section
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_T": 16}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_T": 16}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_T": 32}, num_warps=8, num_stages=2),
+    ],
+    key=["topk_rounded"],
+)
+@triton.jit
+def _tiled_sparse_decode_kernel(
+    # Q: [B, H, D] bf16
+    Q_ptr,
+    # Paged KV cache — three typed views of same underlying memory
+    cache_fp8_ptr,  # float8_e4m3fn flat (1 byte/elem) — for nope
+    cache_uint8_ptr,  # uint8 flat (1 byte/elem) — for scales
+    cache_bf16_ptr,  # bfloat16 flat (2 bytes/elem) — for rope
+    # Indices: [B, topk] int32
+    indices_ptr,
+    # Valid lengths: [B] int32
+    topk_len_ptr,
+    # Output: [B, H, D] bf16 and LSE: [B, H] float32
+    O_ptr,
+    LSE_ptr,
+    # Scalars
+    sm_scale: tl.float32,
+    page_size: tl.int32,
+    page_bytes: tl.int64,
+    scale_section_off: tl.int64,  # page_size * 576
+    H: tl.int32,
+    topk: tl.int32,
+    topk_rounded: tl.int32,  # for autotune key
+    has_topk_len: tl.constexpr,
+    # Strides
+    stride_qb: tl.int32,
+    stride_qh: tl.int32,
+    stride_ob: tl.int32,
+    stride_oh: tl.int32,
+    stride_ib: tl.int32,  # indices batch stride
+    # Constexprs
+    NOPE_PAD: tl.constexpr,  # 512 (padded from 448)
+    ROPE_DIM: tl.constexpr,  # 64
+    NOPE_DIM_RT: tl.int32,  # 448 (runtime, for masking)
+    BLOCK_T: tl.constexpr,  # tokens per tile (16 or 32)
+):
+    """Tiled sparse decode: vectorized gather + QK + softmax + V accumulation.
+
+    Grid: (B, H) — one block per (batch, head) pair.
+    Each block processes all topk tokens in tiles of BLOCK_T.
+    """
+    bid = tl.program_id(0)
+    hid = tl.program_id(1)
+
+    # ---- Load Q for this (batch, head) ----
+    q_base = bid * stride_qb + hid * stride_qh
+    nope_offs = tl.arange(0, NOPE_PAD)  # [512]
+    nope_mask = nope_offs < NOPE_DIM_RT  # [512], True for [0:448]
+    rope_offs = tl.arange(0, ROPE_DIM)  # [64]
+
+    q_nope = tl.load(Q_ptr + q_base + nope_offs, mask=nope_mask, other=0.0)
+    q_nope = q_nope.to(tl.float32) * sm_scale
+    q_rope = tl.load(Q_ptr + q_base + NOPE_DIM_RT + rope_offs)
+    q_rope = q_rope.to(tl.float32) * sm_scale
+
+    # ---- Valid token count ----
+    valid_topk = topk
+    if has_topk_len:
+        valid_topk = tl.load(topk_len_ptr + bid).to(tl.int32)
+        valid_topk = tl.minimum(valid_topk, topk)
+
+    # ---- Online softmax state (base-2 math for SM120 efficiency) ----
+    m_i: tl.float32 = -1e30
+    l_i: tl.float32 = 0.0
+    acc_nope = tl.zeros([NOPE_PAD], dtype=tl.float32)
+    acc_rope = tl.zeros([ROPE_DIM], dtype=tl.float32)
+
+    # ---- Precompute constant index vectors ----
+    group_ids = (nope_offs // 64).to(tl.int64)  # [NOPE_PAD], scale group for each dim
+    t_offs = tl.arange(0, BLOCK_T)  # [BLOCK_T], token offsets within tile
+
+    # ---- Process tokens in tiles of BLOCK_T ----
+    for tile_start in range(0, topk, BLOCK_T):
+        t_idx = tile_start + t_offs  # [BLOCK_T], global token indices
+        t_in_bounds = t_idx < topk  # bounds for index load
+        t_valid = t_idx < valid_topk  # bounds for actual processing
+
+        # Load indices for this tile: [BLOCK_T]
+        raw_indices = tl.load(
+            indices_ptr + bid * stride_ib + t_idx,
+            mask=t_in_bounds,
+            other=-1,
+        )
+        idx_valid = t_valid & (raw_indices >= 0)  # [BLOCK_T] mask
+
+        # Page addressing: [BLOCK_T] (clamp for safe addressing of invalid tokens)
+        safe_indices = tl.where(idx_valid, raw_indices, tl.zeros_like(raw_indices))
+        page_ids = (safe_indices // page_size).to(tl.int64)
+        page_offs_t = (safe_indices % page_size).to(tl.int64)
+        token_data_bases = page_ids * page_bytes + page_offs_t * 576  # [BLOCK_T] int64
+
+        # ---- Vectorized NOPE FP8 gather: [BLOCK_T, NOPE_PAD] ----
+        nope_addrs = token_data_bases[:, None] + nope_offs[None, :].to(tl.int64)
+        nope_2d_mask = idx_valid[:, None] & nope_mask[None, :]
+        kv_nope_fp8 = tl.load(
+            cache_fp8_ptr + nope_addrs,
+            mask=nope_2d_mask,
+            other=0.0,
+        )
+
+        # ---- Vectorized scale gather + dequant: [BLOCK_T, NOPE_PAD] ----
+        scale_bases = page_ids * page_bytes + scale_section_off + page_offs_t * 8
+        scale_addrs = scale_bases[:, None] + group_ids[None, :]
+        scale_raw = tl.load(
+            cache_uint8_ptr + scale_addrs,
+            mask=nope_2d_mask,
+            other=127,
+        )
+        scale_f32 = tl.math.exp2(scale_raw.to(tl.float32) - 127.0)
+        kv_nope = tl.where(nope_2d_mask, kv_nope_fp8.to(tl.float32) * scale_f32, 0.0)
+
+        # ---- Vectorized ROPE BF16 gather: [BLOCK_T, ROPE_DIM] ----
+        rope_byte_bases = token_data_bases + 448
+        rope_elem_bases = (rope_byte_bases // 2).to(tl.int64)
+        rope_addrs = rope_elem_bases[:, None] + rope_offs[None, :].to(tl.int64)
+        kv_rope = tl.load(
+            cache_bf16_ptr + rope_addrs,
+            mask=idx_valid[:, None],
+            other=0.0,
+        ).to(tl.float32)
+
+        # ---- Vectorized QK scores: [BLOCK_T] ----
+        # scores[t] = dot(q_nope, kv_nope[t]) + dot(q_rope, kv_rope[t])
+        scores = tl.sum(q_nope[None, :] * kv_nope, axis=1) + tl.sum(
+            q_rope[None, :] * kv_rope, axis=1
+        )
+        scores = tl.where(idx_valid, scores, -1e30)
+
+        # ---- Online softmax update (base-2, tile-level) ----
+        scores_log2 = scores * LOG2E  # [BLOCK_T]
+        tile_max = tl.max(scores_log2)  # scalar
+        m_new = tl.maximum(m_i, tile_max)
+
+        alpha = tl.math.exp2(m_i - m_new)  # rescale factor
+        p = tl.math.exp2(scores_log2 - m_new)  # [BLOCK_T] attention weights
+        p = tl.where(idx_valid, p, 0.0)  # zero out invalid
+
+        l_i = l_i * alpha + tl.sum(p)
+
+        # ---- Vectorized V accumulation (K=V in MLA) ----
+        # acc += sum_t(p[t] * kv[t, :]) for both nope and rope
+        acc_nope = acc_nope * alpha + tl.sum(p[:, None] * kv_nope, axis=0)
+        acc_rope = acc_rope * alpha + tl.sum(p[:, None] * kv_rope, axis=0)
+        m_i = m_new
+
+    # ---- Normalize output ----
+    safe_l = tl.where(l_i > 0.0, l_i, 1.0)
+    acc_nope = acc_nope / safe_l
+    acc_rope = acc_rope / safe_l
+
+    # LSE: convert from log2 back to natural log
+    lse = tl.where(l_i > 0.0, m_i / LOG2E + tl.math.log(safe_l), float("-inf"))
+
+    # ---- Store output ----
+    o_base = bid * stride_ob + hid * stride_oh
+    tl.store(O_ptr + o_base + nope_offs, acc_nope.to(tl.bfloat16), mask=nope_mask)
+    tl.store(O_ptr + o_base + NOPE_DIM_RT + rope_offs, acc_rope.to(tl.bfloat16))
+    tl.store(LSE_ptr + bid * H + hid, lse)
+
+
+def _run_triton_sparse_decode(
+    q: torch.Tensor,  # [B, 1, H, D] bf16
+    k_cache: torch.Tensor,  # [num_pages, page_size, 1, bpt] float8
+    indices: torch.Tensor,  # [B, ...] int32
+    topk_length: Optional[torch.Tensor],
+    softmax_scale: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Run the tiled Triton sparse decode kernel on one paged KV cache."""
+    B, _, H, D = q.shape
+    num_pages = k_cache.shape[0]
+    page_size = k_cache.shape[1]
+    page_bytes = k_cache.stride(0)  # elements = bytes for float8
+
+    # Flatten indices to [B, topk]
+    flat_indices = indices.reshape(B, -1).contiguous()
+    topk = flat_indices.shape[1]
+
+    # Create three typed views of the flat cache memory.
+    # The KV cache may arrive as uint8 or float8_e4m3fn depending on the
+    # sglang version.  Ensure each view has the correct dtype so Triton
+    # interprets the loaded values correctly (FP8 dequant vs raw integer).
+    total_elems = num_pages * page_bytes
+    raw_flat = k_cache.as_strided((total_elems,), (1,))
+    raw_uint8 = raw_flat.view(torch.uint8)
+    raw_fp8 = raw_uint8.view(torch.float8_e4m3fn)
+    raw_bf16 = raw_uint8.view(torch.bfloat16)
+
+    # Squeeze Q: [B, H, D]
+    q3 = q.squeeze(1)
+    if not q3.is_contiguous():
+        q3 = q3.contiguous()
+
+    out = torch.zeros(B, H, D, dtype=torch.bfloat16, device=q.device)
+    lse = torch.full((B, H), float("-inf"), dtype=torch.float32, device=q.device)
+
+    # Round topk for autotune key stability
+    topk_rounded = triton.next_power_of_2(topk)
+
+    grid = (B, H)
+    _tiled_sparse_decode_kernel[grid](
+        q3,
+        raw_fp8,
+        raw_uint8,
+        raw_bf16,
+        flat_indices,
+        (
+            topk_length
+            if topk_length is not None
+            else torch.empty(0, device=q.device, dtype=torch.int32)
+        ),
+        out,
+        lse,
+        softmax_scale,
+        page_size,
+        int(page_bytes),  # page_bytes (int64)
+        int(page_size * _TOKEN_DATA_STRIDE),  # scale_section_off (int64)
+        H,
+        topk,
+        topk_rounded,
+        topk_length is not None,
+        q3.stride(0),
+        q3.stride(1),
+        out.stride(0),
+        out.stride(1),
+        flat_indices.stride(0),
+        NOPE_PAD=512,
+        ROPE_DIM=_ROPE_DIM,
+        NOPE_DIM_RT=_NOPE_DIM,
+    )
+
+    # Return [B, 1, H, D] and [B, 1, H]
+    return out.unsqueeze(1), lse.unsqueeze(1)
+
+
+def _merge_partial_attn(
+    out1: torch.Tensor,
+    lse1: torch.Tensor,
+    out2: torch.Tensor,
+    lse2: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Merge two attention outputs using LSE-weighted combination.
+
+    out: [B, 1, H, D] bf16,  lse: [B, 1, H] float32
+    """
+    max_lse = torch.maximum(lse1, lse2)
+    w1 = torch.where(lse1 > -1e20, torch.exp(lse1 - max_lse), torch.zeros_like(lse1))
+    w2 = torch.where(lse2 > -1e20, torch.exp(lse2 - max_lse), torch.zeros_like(lse2))
+    total = (w1 + w2).clamp(min=1e-20)
+    merged = (
+        w1.unsqueeze(-1) * out1.float() + w2.unsqueeze(-1) * out2.float()
+    ) / total.unsqueeze(-1)
+    merged_lse = max_lse + torch.log(total)
+    return merged.to(torch.bfloat16), merged_lse
+
+
+def _apply_attn_sink(
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    attn_sink: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply attention sink normalization.
+
+    The sink adds to the softmax denominator without contributing output,
+    effectively down-weighting all attention scores.
+
+    out: [B, 1, H, D] bf16,  lse: [B, 1, H] f32,  attn_sink: [H] f32
+    """
+    sink_lse = attn_sink.view(1, 1, -1).expand_as(lse)
+    combined_lse = torch.logaddexp(lse, sink_lse)
+    w = torch.where(
+        lse > -1e20,
+        torch.exp(lse - combined_lse),
+        torch.zeros_like(lse),
+    )
+    return (out.float() * w.unsqueeze(-1)).to(torch.bfloat16), combined_lse
+
+
+def flash_mla_sparse_decode_triton(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    indices: torch.Tensor,
+    topk_length: Optional[torch.Tensor],
+    attn_sink: Optional[torch.Tensor],
+    head_dim_v: int,
+    softmax_scale: float,
+    extra_k_cache: Optional[torch.Tensor] = None,
+    extra_indices: Optional[torch.Tensor] = None,
+    extra_topk_length: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """SM120-optimized sparse MLA decode using tiled Triton kernel.
+
+    Processes SWA and extra (c4/c128) caches separately via the same
+    Triton kernel, then merges results using LSE-weighted combination.
+    """
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    # Process main cache (SWA)
+    out, lse = _run_triton_sparse_decode(
+        q,
+        k_cache,
+        indices,
+        topk_length,
+        softmax_scale,
+    )
+
+    # Process extra cache (c4 / c128) if present
+    if extra_k_cache is not None and extra_indices is not None:
+        out_extra, lse_extra = _run_triton_sparse_decode(
+            q,
+            extra_k_cache,
+            extra_indices,
+            extra_topk_length,
+            softmax_scale,
+        )
+        out, lse = _merge_partial_attn(out, lse, out_extra, lse_extra)
+
+    # Apply attention sink
+    if attn_sink is not None:
+        out, lse = _apply_attn_sink(out, lse, attn_sink)
+
+    # Return format matching PyTorch fallback: (out, lse.permute(0,2,1))
+    return out, lse.permute(0, 2, 1)

python/sglang/srt/layers/deep_gemm_wrapper/configurer.py

@@ -18,6 +18,9 @@ def _compute_enable_deep_gemm():
     sm_version = get_device_sm()
     if (_is_cuda and sm_version < 90) or (_is_musa and sm_version < 31):
         return False
+    # DeepGEMM requires TMEM/tcgen05 (SM100+datacenter), not available on SM120
+    if sm_version == 120:
+        return False
     if not (_is_cuda or _is_musa):
         return False
 

python/sglang/srt/layers/moe/fused_moe_triton/mxfp4_moe_sm120_triton.py

@@ -0,0 +1,454 @@
+"""SM120-optimized Triton MXFP4 MoE kernel — CUDA graph compatible.
+
+Replaces the PyTorch fallback (per-expert for-loop + full dequant + matmul)
+with fused Triton kernels that:
+1. Fuse FP4 dequant + GEMV (no intermediate BF16 weight materialization)
+2. Process each (token, expert) slot independently — no data-dependent routing
+3. Respect SM120 shared memory constraint (99 KB/block)
+
+CUDA graph compatibility:
+- No .unique(), .item(), .nonzero() — all routing is tensor-level
+- Fixed grid dimensions (M*topk, N_blocks) per captured batch size
+- All control flow is static or within Triton kernels
+
+SM120 constraints:
+- SMEM: 99 KB/block (vs SM100 228 KB)
+- No TMEM/tcgen05 — uses mma.sync.aligned via Triton
+- Max warps: 48/SM
+- Registers: ~128/thread practical limit
+"""
+
+import logging
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _dequant_fp4_lut(nibble):
+    """Decode a 4-bit FP4 E2M1 nibble to float32 using arithmetic."""
+    sign_bit = (nibble >> 3) & 1
+    exp_bits = (nibble >> 1) & 3
+    man_bit = nibble & 1
+
+    is_subnormal = exp_bits == 0
+    mantissa = 1.0 + man_bit.to(tl.float32) * 0.5
+    exponent = tl.math.exp2((exp_bits - 1).to(tl.float32))
+    val = tl.where(is_subnormal, man_bit.to(tl.float32) * 0.5, mantissa * exponent)
+    val = tl.where(sign_bit != 0, -val, val)
+    return val
+
+
+# ── Per-slot GEMV kernel: processes one (token, expert) pair ──
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_N": 32, "BLOCK_K": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 128}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=2),
+    ],
+    key=["N", "K"],
+)
+@triton.jit
+def _mxfp4_slot_gemv_kernel(
+    # Pointers
+    A_ptr,  # [M_total, K] bf16 — source rows
+    B_packed_ptr,  # [E, N, K//2] uint8 — packed FP4 expert weights
+    B_scale_ptr,  # [E, N, K//32] float32 — weight scales
+    C_ptr,  # [num_slots, N] bf16 — output
+    token_ids_ptr,  # [num_slots] int32 — which A row for each slot
+    expert_ids_ptr,  # [num_slots] int32 — which expert's B for each slot
+    # Dimensions
+    N: tl.int32,
+    K: tl.int32,
+    # A strides
+    stride_am: tl.int32,
+    # B strides (within an expert)
+    stride_bn: tl.int32,
+    stride_bk2: tl.int32,
+    # B_scale strides (within an expert)
+    stride_bsn: tl.int32,
+    stride_bsk32: tl.int32,
+    # Expert strides (between experts)
+    expert_b_stride: tl.int64,
+    expert_s_stride: tl.int64,
+    # C strides
+    stride_cm: tl.int32,
+    # Block sizes
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """Per-slot fused MXFP4 dequant + GEMV.
+
+    Grid: (num_slots, cdiv(N, BLOCK_N))
+    Each program computes one (token, expert) pair for a BLOCK_N slice of output.
+    """
+    slot_id = tl.program_id(0)
+    n_block = tl.program_id(1)
+
+    token_id = tl.load(token_ids_ptr + slot_id).to(tl.int64)
+    expert_id = tl.load(expert_ids_ptr + slot_id).to(tl.int64)
+
+    offs_n = n_block * BLOCK_N + tl.arange(0, BLOCK_N)
+    n_mask = offs_n < N
+
+    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
+
+    # Expert weight base pointers
+    b_base = expert_id * expert_b_stride
+    s_base = expert_id * expert_s_stride
+    a_base = token_id * stride_am
+
+    for k_start in range(0, K, BLOCK_K):
+        # ── Load packed B: [BLOCK_N, BLOCK_K//2] ──
+        offs_k2 = k_start // 2 + tl.arange(0, BLOCK_K // 2)
+        b_mask = n_mask[:, None] & (offs_k2[None, :] < K // 2)
+        b_packed = tl.load(
+            B_packed_ptr
+            + b_base
+            + offs_n[:, None] * stride_bn
+            + offs_k2[None, :] * stride_bk2,
+            mask=b_mask,
+            other=0,
+        )
+
+        # ── FP4 dequant ──
+        b_u8 = b_packed.to(tl.int32)
+        val_lo = _dequant_fp4_lut(b_u8 & 0x0F)  # even K indices
+        val_hi = _dequant_fp4_lut((b_u8 >> 4) & 0x0F)  # odd K indices
+
+        # ── Load and apply scales: [BLOCK_N, BLOCK_K//2] ──
+        group_ids = tl.arange(0, BLOCK_K // 2) // 16  # 32 values per group, 2 per byte
+        s_mask = n_mask[:, None] & ((k_start // 32 + group_ids[None, :]) < K // 32)
+        scales = tl.load(
+            B_scale_ptr
+            + s_base
+            + offs_n[:, None] * stride_bsn
+            + (k_start // 32 + group_ids[None, :]) * stride_bsk32,
+            mask=s_mask,
+            other=1.0,
+        )
+        val_lo = val_lo * scales
+        val_hi = val_hi * scales
+
+        # ── Load A even/odd: [BLOCK_K//2] each ──
+        offs_k_even = k_start + tl.arange(0, BLOCK_K // 2) * 2
+        offs_k_odd = offs_k_even + 1
+
+        a_even = tl.load(
+            A_ptr + a_base + offs_k_even,
+            mask=offs_k_even < K,
+            other=0.0,
+        ).to(tl.float32)
+        a_odd = tl.load(
+            A_ptr + a_base + offs_k_odd,
+            mask=offs_k_odd < K,
+            other=0.0,
+        ).to(tl.float32)
+
+        # ── Dot product: acc[n] += sum_k(a_even[k]*B_lo[n,k] + a_odd[k]*B_hi[n,k]) ──
+        acc += tl.sum(a_even[None, :] * val_lo, axis=1)
+        acc += tl.sum(a_odd[None, :] * val_hi, axis=1)
+
+    # ── Store output ──
+    tl.store(
+        C_ptr + slot_id * stride_cm + offs_n,
+        acc.to(tl.bfloat16),
+        mask=n_mask,
+    )
+
+
+# ── Legacy per-expert GEMM kernel (kept for benchmarking) ──
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 64}, num_warps=4, num_stages=2
+        ),
+        triton.Config(
+            {"BLOCK_M": 32, "BLOCK_N": 32, "BLOCK_K": 64}, num_warps=4, num_stages=2
+        ),
+        triton.Config(
+            {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64}, num_warps=4, num_stages=2
+        ),
+        triton.Config(
+            {"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_warps=8, num_stages=2
+        ),
+        triton.Config(
+            {"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 64}, num_warps=4, num_stages=2
+        ),
+    ],
+    key=["M", "N", "K"],
+)
+@triton.jit
+def _mxfp4_gemm_kernel(
+    # Pointers
+    A_ptr,  # [M, K] bf16 activation
+    B_packed_ptr,  # [N, K//2] uint8 packed FP4
+    B_scale_ptr,  # [N, K//32] float32 scales
+    C_ptr,  # [M, N] bf16 output
+    # Dimensions
+    M,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bn,
+    stride_bk2,
+    stride_bsn,
+    stride_bsk32,
+    stride_cm,
+    stride_cn,
+    # Constexprs
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """Fused MXFP4 dequant + GEMM: C = A @ dequant(B_packed, B_scale).T"""
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k_start in range(0, K, BLOCK_K):
+        offs_k2 = k_start // 2 + tl.arange(0, BLOCK_K // 2)
+        b_mask = (offs_n[:, None] < N) & (offs_k2[None, :] < K // 2)
+        b_packed = tl.load(
+            B_packed_ptr + offs_n[:, None] * stride_bn + offs_k2[None, :] * stride_bk2,
+            mask=b_mask,
+            other=0,
+        )
+
+        b_u8 = b_packed.to(tl.int32)
+        val_lo = _dequant_fp4_lut(b_u8 & 0x0F)
+        val_hi = _dequant_fp4_lut((b_u8 >> 4) & 0x0F)
+
+        group_ids = tl.arange(0, BLOCK_K // 2) // 16
+        scales_per_byte = tl.load(
+            B_scale_ptr
+            + offs_n[:, None] * stride_bsn
+            + (k_start // 32 + group_ids[None, :]) * stride_bsk32,
+            mask=(offs_n[:, None] < N)
+            & ((k_start // 32 + group_ids[None, :]) < K // 32),
+            other=1.0,
+        )
+        val_lo = val_lo * scales_per_byte
+        val_hi = val_hi * scales_per_byte
+
+        offs_k_even = k_start + tl.arange(0, BLOCK_K // 2) * 2
+        offs_k_odd = offs_k_even + 1
+
+        a_even_mask = (offs_m[:, None] < M) & (offs_k_even[None, :] < K)
+        a_even = tl.load(
+            A_ptr + offs_m[:, None] * stride_am + offs_k_even[None, :] * stride_ak,
+            mask=a_even_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        a_odd_mask = (offs_m[:, None] < M) & (offs_k_odd[None, :] < K)
+        a_odd = tl.load(
+            A_ptr + offs_m[:, None] * stride_am + offs_k_odd[None, :] * stride_ak,
+            mask=a_odd_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        acc += tl.dot(a_even, tl.trans(val_lo))
+        acc += tl.dot(a_odd, tl.trans(val_hi))
+
+    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    tl.store(
+        C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn,
+        acc.to(tl.bfloat16),
+        mask=c_mask,
+    )
+
+
+def mxfp4_gemm_triton(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_scale: torch.Tensor,
+    K_full: int,
+) -> torch.Tensor:
+    """Triton fused MXFP4 dequant + GEMM: output = A @ dequant(B).T
+
+    Kept for standalone benchmarking. The MoE forward uses the slot kernel.
+    """
+    M = A.shape[0]
+    N = B_packed.shape[0]
+    K = K_full
+
+    if B_scale.dtype == torch.float8_e8m0fnu:
+        B_scale = B_scale.to(torch.float32)
+    elif B_scale.dtype != torch.float32:
+        B_scale = B_scale.float()
+
+    C = torch.empty(M, N, dtype=torch.bfloat16, device=A.device)
+    A = A.contiguous()
+    B_packed = B_packed.contiguous()
+    B_scale = B_scale.contiguous()
+
+    grid = lambda meta: (
+        triton.cdiv(M, meta["BLOCK_M"]),
+        triton.cdiv(N, meta["BLOCK_N"]),
+    )
+    B_u8 = B_packed.view(torch.uint8)
+
+    _mxfp4_gemm_kernel[grid](
+        A,
+        B_u8,
+        B_scale,
+        C,
+        M,
+        N,
+        K,
+        A.stride(0),
+        A.stride(1),
+        B_u8.stride(0),
+        B_u8.stride(1),
+        B_scale.stride(0),
+        B_scale.stride(1),
+        C.stride(0),
+        C.stride(1),
+    )
+    return C
+
+
+def mxfp4_moe_forward_triton(
+    hidden_states: torch.Tensor,
+    w13_packed: torch.Tensor,
+    w2_packed: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    hidden_size: int,
+    intermediate_size: int,
+    inplace: bool = False,
+    routed_scaling_factor: Optional[float] = None,
+    clamp_limit: Optional[float] = None,
+) -> torch.Tensor:
+    """SM120-optimized MXFP4 MoE forward — CUDA graph compatible.
+
+    Uses per-slot GEMV kernels instead of per-expert Python loops.
+    Each (token, expert) slot is processed independently with a fixed grid,
+    eliminating .unique()/.item()/.nonzero() that break CUDA graph capture.
+    """
+    import torch.nn.functional as F
+
+    M, K = hidden_states.shape
+    topk = topk_ids.shape[1]
+    I = intermediate_size
+    num_slots = M * topk
+    device = hidden_states.device
+    dtype = hidden_states.dtype
+
+    # ── Graph-safe routing: flatten topk assignments ──
+    # token_ids[slot] = which row of A (original token index)
+    # expert_ids[slot] = which expert's weights to use
+    # topk_ids may contain -1 for padded/filtered tokens; clamp to 0 for safe
+    # Triton loads, then zero out invalid slots' output after GEMM.
+    flat_expert_ids_raw = topk_ids.reshape(-1).contiguous()  # [M*topk]
+    invalid_slot_mask = flat_expert_ids_raw < 0  # [M*topk]
+    flat_expert_ids = flat_expert_ids_raw.clamp(min=0)  # safe for indexing
+    token_ids = (
+        torch.arange(M, device=device, dtype=torch.int32)
+        .unsqueeze(1)
+        .expand(M, topk)
+        .reshape(-1)
+        .contiguous()
+    )  # [M*topk]
+
+    # ── Ensure scales are float32 ──
+    if w13_scale.dtype != torch.float32:
+        w13_scale = w13_scale.to(torch.float32)
+    if w2_scale.dtype != torch.float32:
+        w2_scale = w2_scale.to(torch.float32)
+
+    # ── GEMM1: gate_up projection ──
+    # hidden_states[token] @ w13[expert].T → [num_slots, 2*I]
+    intermediate = torch.empty(num_slots, 2 * I, dtype=dtype, device=device)
+
+    w13_u8 = w13_packed.view(torch.uint8)  # [E, 2*I, K//2]
+    grid1 = lambda meta: (num_slots, triton.cdiv(2 * I, meta["BLOCK_N"]))
+
+    _mxfp4_slot_gemv_kernel[grid1](
+        hidden_states,
+        w13_u8,
+        w13_scale,
+        intermediate,
+        token_ids,
+        flat_expert_ids,
+        2 * I,
+        K,
+        hidden_states.stride(0),
+        w13_u8.stride(1),
+        w13_u8.stride(2),
+        w13_scale.stride(1),
+        w13_scale.stride(2),
+        w13_u8.stride(0),
+        w13_scale.stride(0),
+        intermediate.stride(0),
+    )
+
+    # ── SiLU activation (graph-safe vectorized ops) ──
+    gate = intermediate[:, :I].float()
+    up = intermediate[:, I:].float()
+    if clamp_limit is not None and clamp_limit > 0:
+        gate = torch.clamp(gate, max=clamp_limit)
+        up = torch.clamp(up, min=-clamp_limit, max=clamp_limit)
+    activated = (F.silu(gate) * up).to(dtype)
+
+    # ── GEMM2: down projection ──
+    # activated[slot] @ w2[expert].T → [num_slots, K]
+    down = torch.empty(num_slots, K, dtype=dtype, device=device)
+
+    # For GEMM2, A is the activated buffer — each slot reads its own row
+    slot_ids = torch.arange(num_slots, device=device, dtype=torch.int32)
+
+    w2_u8 = w2_packed.view(torch.uint8)  # [E, K, I//2]
+    grid2 = lambda meta: (num_slots, triton.cdiv(K, meta["BLOCK_N"]))
+
+    _mxfp4_slot_gemv_kernel[grid2](
+        activated,
+        w2_u8,
+        w2_scale,
+        down,
+        slot_ids,
+        flat_expert_ids,
+        K,
+        I,
+        activated.stride(0),
+        w2_u8.stride(1),
+        w2_u8.stride(2),
+        w2_scale.stride(1),
+        w2_scale.stride(2),
+        w2_u8.stride(0),
+        w2_scale.stride(0),
+        down.stride(0),
+    )
+
+    # ── Zero out invalid slots (padded/filtered tokens with topk_ids == -1) ──
+    # Use multiplication instead of boolean indexing to stay CUDA-graph-safe
+    # (no GPU→CPU sync). valid_mask is 1.0 for valid slots, 0.0 for invalid.
+    valid_mask = (~invalid_slot_mask).unsqueeze(1).to(dtype)  # [M*topk, 1]
+    down = down * valid_mask
+
+    # ── Weighted sum across topk slots (graph-safe) ──
+    flat_weights = topk_weights.reshape(-1).unsqueeze(1).to(dtype)  # [M*topk, 1]
+    output = (down * flat_weights).view(M, topk, K).sum(dim=1)
+
+    if routed_scaling_factor is not None and routed_scaling_factor != 1.0:
+        output.mul_(routed_scaling_factor)
+
+    return output

python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py

@@ -9,7 +9,7 @@ from torch.nn import Module
 from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
 from sglang.srt.layers.moe.utils import MoeRunnerBackend
 from sglang.srt.utils import log_info_on_rank0, set_weight_attrs
-from sglang.srt.utils.common import is_sm90_supported
+from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
@@ -42,6 +42,7 @@ class Mxfp4MarlinMoEMethod:
             FusedMoeWeightScaleSupported,
         )
 
+        layer._dsv4_mxfp4_backend = None  # set in process_weights_after_loading
         fp4_block_k = 32
 
         w13_weight = torch.nn.Parameter(
@@ -108,10 +109,42 @@ class Mxfp4MarlinMoEMethod:
         if getattr(layer, "_mega_moe_weights_built", False):
             return
 
-        if not is_sm90_supported():
+        if not is_sm90_supported() and not is_sm120_supported():
             raise RuntimeError(
                 "DeepSeekV4 MXFP4 Marlin fallback requires Hopper/SM90 or above."
             )
+
+        # SM120: Skip Marlin repacking, keep original weight format
+        # for Triton dequant kernel (Marlin kernel produces NaN on SM120)
+        if is_sm120_supported():
+            from torch.nn import Parameter
+
+            log_info_on_rank0(
+                logger,
+                f"SM120 detected: using PyTorch MXFP4 MoE fallback "
+                f"(layer: {self.prefix})...",
+            )
+            # Keep weights in original packed int8 format
+            # Normalize scales to float32 for direct use in dequant
+            w13_s = layer.w13_weight_scale_inv.data
+            w2_s = layer.w2_weight_scale_inv.data
+            if w13_s.dtype == torch.float8_e8m0fnu:
+                pass  # already in e8m0 format, will convert at runtime
+            elif w13_s.dtype in (torch.uint8, torch.int8):
+                layer.w13_weight_scale_inv = Parameter(
+                    w13_s.view(torch.uint8)
+                    .view(torch.float8_e8m0fnu)
+                    .to(torch.float32),
+                    requires_grad=False,
+                )
+                layer.w2_weight_scale_inv = Parameter(
+                    w2_s.view(torch.uint8).view(torch.float8_e8m0fnu).to(torch.float32),
+                    requires_grad=False,
+                )
+            # else: float32 scales are already usable directly
+            layer._dsv4_mxfp4_backend = "sm120_triton"
+            return
+
         if not check_moe_marlin_supports_layer(layer, 32):
             raise RuntimeError(
                 "Current DeepSeekV4 MoE layer does not satisfy Marlin constraints."
@@ -144,6 +177,43 @@ class Mxfp4MarlinMoEMethod:
         if not TopKOutputChecker.format_is_standard(topk_output):
             raise ValueError(f"Unsupported topk output format: {topk_output.format}")
 
+        # SM120: use Triton fused dequant+GEMM (Marlin kernel produces NaN on SM120)
+        if layer._dsv4_mxfp4_backend == "sm120_triton":
+            from sglang.srt.layers.moe.fused_moe_triton.mxfp4_moe_sm120_triton import (
+                mxfp4_moe_forward_triton,
+            )
+
+            hidden_states = dispatch_output.hidden_states
+            w13 = layer.w13_weight.data
+            w2 = layer.w2_weight.data
+            w13_scale = layer.w13_weight_scale_inv.data
+            w2_scale = layer.w2_weight_scale_inv.data
+            intermediate_size = w13.shape[1] // 2
+            hidden_size = w13.shape[2] * 2
+
+            output = mxfp4_moe_forward_triton(
+                hidden_states=hidden_states,
+                w13_packed=w13,
+                w2_packed=w2,
+                w13_scale=w13_scale,
+                w2_scale=w2_scale,
+                topk_ids=topk_output.topk_ids,
+                topk_weights=topk_output.topk_weights,
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                routed_scaling_factor=(
+                    self.runner.config.routed_scaling_factor
+                    if hasattr(self.runner, "config")
+                    else None
+                ),
+                clamp_limit=(
+                    self.runner.config.swiglu_limit
+                    if hasattr(self.runner, "config")
+                    else None
+                ),
+            )
+            return StandardCombineInput(hidden_states=output)
+
         quant_info = MarlinMoeQuantInfo(
             w13_qweight=layer.w13_weight,
             w2_qweight=layer.w2_weight,

python/sglang/srt/server_args.py

@@ -2066,6 +2066,20 @@ class ServerArgs:
 
             validate_deepseek_v4_cp(self)
 
+            if is_sm120_supported():
+                if self.moe_runner_backend == "auto":
+                    self.moe_runner_backend = "marlin"
+                    logger.info(
+                        "Use marlin as MoE runner backend on SM120 for DeepseekV4"
+                    )
+                # SM120 lacks tcgen05/TMEM: disable features that depend on
+                # DeepGEMM or require >99KB SMEM (topk_v2).
+                envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
+                envs.SGLANG_OPT_USE_TOPK_V2.set(False)
+                envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
+                envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
+                envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.set(True)
+
         elif model_arch in ["GptOssForCausalLM"]:
             # Set attention backend for GPT-OSS
             if self.is_attention_backend_not_set():

test/registered/kernels/test_sm120_flash_mla.py

@@ -0,0 +1,465 @@
+"""SM120 FlashMLA sparse decode unit tests.
+
+Validates the SM120-specific FlashMLA implementation that replaces the upstream
+`flash_mla` CUDA kernel (unavailable on SM120 / RTX PRO 6000):
+
+- ``_gather_and_dequant``: byte-precise dequant of paged FP8 + BF16 + UE8M0 KV
+  cache. Covers the dtype-reinterpretation surface that caused the historic
+  uint8 garbled-output regression (see progress doc §4).
+- ``_sm120_sparse_decode_fwd``: pure-PyTorch reference path.
+- ``flash_mla_sparse_decode_triton``: tiled Triton kernel.
+- ``_apply_attn_sink`` / ``_merge_partial_attn``: post-processing helpers.
+- ``flash_mla_with_kvcache_sm120``: entry-point dispatch on
+  ``SGLANG_SM120_TRITON_FLASHMLA`` selects torch/triton paths and both yield
+  matching output.
+
+DSv4 cache layout (per page):
+    data section:  page_size * 576 bytes   = 64 tokens * (448 nope + 128 rope)
+    scale section: page_size * 8 bytes     = 64 tokens * (7 UE8M0 scales + 1 pad)
+    total bytes:   page_size * 584         = stride(0) of k_cache
+"""
+
+from __future__ import annotations
+
+import unittest
+from unittest import mock
+
+import torch
+
+from sglang.srt.layers.attention import flash_mla_sm120 as fmod
+from sglang.srt.layers.attention.flash_mla_sm120 import (
+    _D,
+    _NOPE_DIM,
+    _NOPE_ROPE_STRIDE,
+    _NUM_TILES,
+    _ROPE_DIM,
+    _SCALE_STRIDE,
+    _TILE_SIZE,
+    _gather_and_dequant,
+    _sm120_sparse_decode_fwd,
+    flash_mla_with_kvcache_sm120,
+)
+from sglang.srt.layers.attention.flash_mla_sm120_triton import (
+    _apply_attn_sink,
+    _merge_partial_attn,
+    flash_mla_sparse_decode_triton,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b", runner_config="1-gpu-large")
+
+
+# Per-token byte layout
+_BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE  # 576 + 8 = 584
+
+
+def _build_kvcache(
+    num_pages: int,
+    page_size: int,
+    *,
+    device: torch.device,
+    seed: int = 0,
+):
+    """Build a synthetic FP8/BF16/UE8M0 KV cache.
+
+    Returns the (num_pages, page_size, 1, bpt) FP8-viewed cache plus the raw
+    nope FP8 (float), rope BF16 (float), and UE8M0 scale (float) reference
+    tensors for verification.
+    """
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    bpt = _BYTES_PER_TOKEN  # 584 — also satisfies k_cache.stride(0) requirement
+    raw = torch.zeros(num_pages, page_size, bpt, dtype=torch.uint8, device=device)
+
+    # ---- Nope FP8 region: per-token bytes [0:448] ----
+    # Stay in [0, 0x6F] to avoid NaN/Inf in float8_e4m3fn.
+    nope_bytes = torch.randint(
+        0, 0x70, (num_pages, page_size, _NOPE_DIM), generator=g, dtype=torch.uint8
+    ).to(device)
+    raw[:, :, :_NOPE_DIM] = nope_bytes
+
+    # ---- Rope BF16 region: per-token bytes [448:576] = 64 bf16 values ----
+    rope_bf16_vals = (
+        torch.randn((num_pages, page_size, _ROPE_DIM), generator=g, dtype=torch.float32)
+        .clamp(-2.0, 2.0)
+        .to(torch.bfloat16)
+        .to(device)
+    )
+    # View bf16 as 2 bytes per value -> 128 bytes per token rope region
+    rope_as_uint8 = rope_bf16_vals.contiguous().view(
+        torch.uint8
+    )  # (num_pages, page_size, 128)
+    raw[:, :, _NOPE_DIM : _NOPE_DIM + _ROPE_DIM * 2] = rope_as_uint8
+
+    # ---- Scale section: starts at page_size * 576 ----
+    # 7 UE8M0 bytes per token + 1 pad.  Keep exponents in a sane range
+    # (UE8M0 byte 'b' decodes to 2**(b-127)), so 120..130 gives ~[1/128, 8].
+    scale_bytes = torch.randint(
+        120,
+        131,
+        (num_pages, page_size, _NUM_TILES),
+        generator=g,
+        dtype=torch.uint8,
+    ).to(device)
+    # raw is (num_pages, page_size, bpt) with bpt=584; data ends at 576, scale region 576..584
+    # Per-token scale offset relative to scale section: token_idx * 8.
+    # Note: the scale region in raw memory is at the END of the page after the
+    # data section. Since raw shape is (num_pages, page_size, bpt), where
+    # bpt=584, scale lives in raw[:, :, _NOPE_ROPE_STRIDE : _NOPE_ROPE_STRIDE+8].
+    # BUT the gather code expects layout
+    #   raw_pages: (num_pages, page_bytes)
+    #   scale_section_offset = page_size * 576
+    # i.e. data for ALL tokens first, then scale for ALL tokens. So we need
+    # the flat per-page layout: [tok0_data(576), tok1_data, ..., tokN-1_data,
+    # tok0_scale(8), tok1_scale, ...].
+    # We constructed raw as (num_pages, page_size, bpt) which interleaves data
+    # and scales per token.  Build a fresh buffer with the correct flat order.
+    flat = torch.zeros(num_pages, page_size * bpt, dtype=torch.uint8, device=device)
+    # Data: per-token 576 bytes contiguous
+    flat_data = raw[:, :, :_NOPE_ROPE_STRIDE].reshape(
+        num_pages, page_size * _NOPE_ROPE_STRIDE
+    )
+    flat[:, : page_size * _NOPE_ROPE_STRIDE] = flat_data
+    # Scales: per-token 8 bytes (only first 7 written)
+    scale_block = torch.zeros(
+        num_pages, page_size, _SCALE_STRIDE, dtype=torch.uint8, device=device
+    )
+    scale_block[:, :, :_NUM_TILES] = scale_bytes
+    flat[:, page_size * _NOPE_ROPE_STRIDE :] = scale_block.reshape(
+        num_pages, page_size * _SCALE_STRIDE
+    )
+
+    # View as (num_pages, page_size, 1, bpt) float8_e4m3fn
+    k_cache = flat.view(num_pages, page_size, 1, bpt).view(torch.float8_e4m3fn)
+    # The Triton kernel expects k_cache.stride(0) == page_size * bpt.
+    assert k_cache.stride(0) == page_size * bpt
+
+    # Reference dequant per token (matches what _gather_and_dequant should produce):
+    nope_fp8 = nope_bytes.view(
+        torch.float8_e4m3fn
+    ).float()  # (num_pages, page_size, 448)
+    scale_e8m0 = scale_bytes.view(
+        torch.float8_e8m0fnu
+    ).float()  # (num_pages, page_size, 7)
+    nope_dequant = (
+        nope_fp8.view(num_pages, page_size, _NUM_TILES, _TILE_SIZE)
+        * scale_e8m0.view(num_pages, page_size, _NUM_TILES, 1)
+    ).view(
+        num_pages, page_size, _NOPE_DIM
+    )  # float32
+    ref_per_token = torch.cat(
+        [nope_dequant.to(torch.bfloat16), rope_bf16_vals], dim=-1
+    )  # (num_pages, page_size, 512) bf16
+    return k_cache, ref_per_token
+
+
+def _build_q_indices(
+    batch_size: int,
+    num_heads: int,
+    topk: int,
+    num_pages: int,
+    page_size: int,
+    *,
+    device: torch.device,
+    seed: int = 1,
+):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    q = (
+        torch.randn((batch_size, 1, num_heads, _D), generator=g, dtype=torch.float32)
+        .clamp(-1.5, 1.5)
+        .to(torch.bfloat16)
+        .to(device)
+    )
+    # Each batch picks `topk` random valid token-level indices into the pool
+    pool_size = num_pages * page_size
+    indices = torch.zeros((batch_size, 1, topk), dtype=torch.int32, device=device)
+    for b in range(batch_size):
+        perm = torch.randperm(pool_size, generator=g)[:topk]
+        indices[b, 0] = perm.to(device=device, dtype=torch.int32)
+    return q, indices
+
+
+class TestGatherAndDequant(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def test_basic_dequant_matches_manual(self):
+        """Per-token output equals nope_fp8 * scale concatenated with rope_bf16."""
+        num_pages, page_size = 3, 64
+        k_cache, ref_per_token = _build_kvcache(
+            num_pages, page_size, device=self.device, seed=42
+        )
+        # Pick 1 index per page (positions 0, 32, 63)
+        token_ids = torch.tensor(
+            [[0, 1 * page_size + 32, 2 * page_size + 63]],
+            dtype=torch.int32,
+            device=self.device,
+        )
+        out = _gather_and_dequant(k_cache, token_ids, page_size)
+        self.assertEqual(out.shape, (1, 3, _D))
+        # Expected entries
+        expected = torch.stack(
+            [
+                ref_per_token[0, 0],
+                ref_per_token[1, 32],
+                ref_per_token[2, 63],
+            ],
+            dim=0,
+        ).unsqueeze(0)
+        # bf16 dequant: allow up to 1 ULP per element.  Scales are integral
+        # powers of 2 so the only loss is the fp8 -> bf16 mantissa rounding.
+        torch.testing.assert_close(out, expected, atol=1e-2, rtol=1e-2)
+
+    def test_dequant_handles_full_page_range(self):
+        """All token positions in a single page produce correct output."""
+        num_pages, page_size = 2, 64
+        k_cache, ref_per_token = _build_kvcache(
+            num_pages, page_size, device=self.device, seed=7
+        )
+        token_ids = torch.arange(page_size, dtype=torch.int32, device=self.device).view(
+            1, page_size
+        )
+        out = _gather_and_dequant(k_cache, token_ids, page_size)
+        torch.testing.assert_close(
+            out, ref_per_token[0].unsqueeze(0), atol=1e-2, rtol=1e-2
+        )
+
+
+class TestSparseDecodeTritonVsTorch(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def _run(
+        self,
+        batch_size: int = 2,
+        num_heads: int = 4,
+        topk: int = 64,
+        page_size: int = 64,
+        num_pages: int = 4,
+        seed: int = 11,
+        with_topk_length: bool = True,
+        with_attn_sink: bool = False,
+    ):
+        k_cache, _ = _build_kvcache(num_pages, page_size, device=self.device, seed=seed)
+        q, indices = _build_q_indices(
+            batch_size,
+            num_heads,
+            topk,
+            num_pages,
+            page_size,
+            device=self.device,
+            seed=seed + 100,
+        )
+        topk_length = None
+        if with_topk_length:
+            topk_length = torch.tensor(
+                [topk // 2, topk] if batch_size == 2 else [topk] * batch_size,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        attn_sink = None
+        if with_attn_sink:
+            attn_sink = torch.full(
+                (num_heads,), -1.5, dtype=torch.float32, device=self.device
+            )
+        softmax_scale = _D ** (-0.5)
+        # Production passes config.v_head_dim == _D (512). Both impls return
+        # (B, 1, H, _D); the Triton kernel always computes the full nope+rope.
+        head_dim_v = _D
+
+        # PyTorch reference
+        ref_out, ref_lse = _sm120_sparse_decode_fwd(
+            q,
+            k_cache,
+            indices,
+            topk_length,
+            attn_sink,
+            head_dim_v=head_dim_v,
+            softmax_scale=softmax_scale,
+        )
+
+        # Triton
+        tri_out, tri_lse = flash_mla_sparse_decode_triton(
+            q,
+            k_cache,
+            indices,
+            topk_length,
+            attn_sink,
+            head_dim_v=head_dim_v,
+            softmax_scale=softmax_scale,
+        )
+
+        # Outputs both bf16, allow drift from online softmax base-2 vs base-e
+        # math + bf16 mantissa precision (~7 bits).  5e-2 abs/rel is generous
+        # but accounts for accumulated rounding across topk tokens.
+        torch.testing.assert_close(
+            tri_out.to(torch.float32),
+            ref_out.to(torch.float32),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+        return ref_out, tri_out, ref_lse, tri_lse
+
+    def test_triton_vs_torch_basic(self):
+        self._run()
+
+    def test_triton_vs_torch_no_topk_length(self):
+        self._run(with_topk_length=False)
+
+    def test_triton_vs_torch_with_attn_sink(self):
+        self._run(with_attn_sink=True)
+
+    def test_triton_vs_torch_small_topk(self):
+        self._run(topk=32, num_heads=2)
+
+    def test_triton_vs_torch_negative_indices_are_masked(self):
+        """Indices < 0 are 'invalid' and must contribute zero to output."""
+        k_cache, _ = _build_kvcache(4, 64, device=self.device, seed=3)
+        q, indices = _build_q_indices(2, 4, 32, 4, 64, device=self.device, seed=99)
+        # Half of each batch's indices set to -1
+        indices[:, :, 16:] = -1
+        topk_length = torch.tensor([32, 32], dtype=torch.int32, device=self.device)
+
+        ref_out, _ = _sm120_sparse_decode_fwd(
+            q, k_cache, indices, topk_length, None, _D, _D**-0.5
+        )
+        tri_out, _ = flash_mla_sparse_decode_triton(
+            q, k_cache, indices, topk_length, None, _D, _D**-0.5
+        )
+        torch.testing.assert_close(
+            tri_out.to(torch.float32),
+            ref_out.to(torch.float32),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+
+class TestApplyAttnSink(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def test_sink_zero_means_full_dampen_to_half(self):
+        """attn_sink == lse implies sink contributes equal weight: output halves."""
+        B, H, D = 1, 2, 8
+        lse = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
+        attn_sink = torch.zeros(H, device=self.device, dtype=torch.float32)
+        out = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
+        new_out, new_lse = _apply_attn_sink(out, lse, attn_sink)
+        # combined_lse = logaddexp(0, 0) = log(2). w = exp(-log2) = 0.5.
+        torch.testing.assert_close(
+            new_out.float(),
+            torch.full_like(new_out.float(), 0.5),
+            atol=2e-3,
+            rtol=2e-3,
+        )
+        torch.testing.assert_close(
+            new_lse,
+            torch.full_like(new_lse, 0.6931472),  # log(2)
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    def test_sink_dead_lse_stays_zero(self):
+        """lse == -inf (no valid tokens) -> output stays zero (weight=0)."""
+        B, H, D = 1, 2, 4
+        lse = torch.full(
+            (B, 1, H), float("-inf"), device=self.device, dtype=torch.float32
+        )
+        attn_sink = torch.zeros(H, device=self.device, dtype=torch.float32)
+        out = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
+        new_out, _ = _apply_attn_sink(out, lse, attn_sink)
+        torch.testing.assert_close(
+            new_out.float(), torch.zeros_like(new_out.float()), atol=0, rtol=0
+        )
+
+
+class TestMergePartialAttn(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def test_equal_lse_arithmetic_mean(self):
+        """lse1 == lse2 -> merged output is the arithmetic mean of out1, out2."""
+        B, H, D = 1, 2, 4
+        lse = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
+        out1 = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
+        out2 = torch.full((B, 1, H, D), 3.0, device=self.device, dtype=torch.bfloat16)
+        merged, merged_lse = _merge_partial_attn(out1, lse, out2, lse)
+        torch.testing.assert_close(
+            merged.float(),
+            torch.full_like(merged.float(), 2.0),
+            atol=1e-2,
+            rtol=1e-2,
+        )
+        torch.testing.assert_close(
+            merged_lse,
+            torch.full_like(merged_lse, 0.6931472),  # log(2)
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    def test_one_dead_branch_passes_through(self):
+        """If lse2 == -inf, merged equals out1."""
+        B, H, D = 1, 1, 4
+        lse1 = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
+        lse2 = torch.full(
+            (B, 1, H), float("-inf"), device=self.device, dtype=torch.float32
+        )
+        out1 = torch.full((B, 1, H, D), 2.5, device=self.device, dtype=torch.bfloat16)
+        out2 = torch.full((B, 1, H, D), 99.0, device=self.device, dtype=torch.bfloat16)
+        merged, merged_lse = _merge_partial_attn(out1, lse1, out2, lse2)
+        torch.testing.assert_close(merged.float(), out1.float(), atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(merged_lse, lse1, atol=1e-5, rtol=1e-5)
+
+
+class TestEntryPointDispatch(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def test_torch_backend_matches_triton_backend(self):
+        """SGLANG_SM120_TRITON_FLASHMLA toggles backend; both return matching out."""
+        k_cache, _ = _build_kvcache(4, 64, device=self.device, seed=5)
+        q, indices = _build_q_indices(1, 4, 32, 4, 64, device=self.device, seed=13)
+        topk_length = torch.tensor([32], dtype=torch.int32, device=self.device)
+
+        kwargs = dict(
+            q=q,
+            k_cache=k_cache,
+            indices=indices,
+            topk_length=topk_length,
+            attn_sink=None,
+            head_dim_v=_D,
+            softmax_scale=_D**-0.5,
+        )
+
+        with mock.patch.object(fmod, "_sm120_default_backend", "torch"):
+            out_torch, _ = flash_mla_with_kvcache_sm120(**kwargs)
+        with mock.patch.object(fmod, "_sm120_default_backend", "triton"):
+            out_triton, _ = flash_mla_with_kvcache_sm120(**kwargs)
+
+        torch.testing.assert_close(
+            out_torch.to(torch.float32),
+            out_triton.to(torch.float32),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(unittest.main())

test/registered/kernels/test_sm120_paged_mqa_logits.py

@@ -0,0 +1,314 @@
+"""SM120 fp8_paged_mqa_logits_torch_sm120 vectorized PyTorch fallback tests.
+
+Validates that the vectorized SM120-specific implementation matches the loopy
+reference (`fp8_paged_mqa_logits_torch`) and is CUDA-graph compatible.
+
+Coverage:
+- Numeric equivalence vs reference at small shapes
+- Both KV-cache dtype views (uint8 raw / float8_e4m3fn) — guards against the
+  historic garbled-output bug where Triton kernels treated uint8 bytes as raw
+  integers instead of FP8 (dsv4_sm120_progress.md §4)
+- Variable per-batch seq_lens with -inf masking semantics
+- CUDA graph capture + replay equivalence (no .item() / data-dependent shapes)
+- Shape-assertion guards
+
+Pure-PyTorch implementations on any CUDA GPU — no SM120 hardware required.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.dsv4.indexer import (
+    FP8_DTYPE,
+    fp8_paged_mqa_logits_torch,
+    fp8_paged_mqa_logits_torch_sm120,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
+
+
+# DSv4 indexer cache layout (fixed by deepseek_v4_memory_pool.DeepSeekV4IndexerPool):
+#   page_size = 64 tokens
+#   head_dim = 128 (FP8 values per token)
+#   quant_block_size = 128 -> num_scales_per_token = 1 (fp32 scale)
+#   per-page memory: [page_size*head_dim FP8 bytes][page_size*4 scale bytes]
+#                  = [8192][256] = 8448 bytes
+# The (block_size, 1, head_dim+4) shape is a fiction for downstream consumers.
+PAGE_SIZE = 64
+HEAD_DIM = 128
+SCALE_BYTES_PER_TOKEN = 4
+HEAD_DIM_WITH_SF = HEAD_DIM + SCALE_BYTES_PER_TOKEN  # 132
+PAGE_BYTES = PAGE_SIZE * HEAD_DIM + PAGE_SIZE * SCALE_BYTES_PER_TOKEN  # 8448
+
+
+def _build_kvcache(
+    num_pages: int,
+    *,
+    dtype_view: torch.dtype,
+    device: torch.device,
+    seed: int = 0,
+) -> torch.Tensor:
+    """Construct a paged KV cache matching the production layout.
+
+    Returns a tensor shaped (num_pages, PAGE_SIZE, 1, HEAD_DIM_WITH_SF) whose
+    underlying memory is the same regardless of `dtype_view`:
+      - [0 : PAGE_SIZE*HEAD_DIM)        : random FP8 bit patterns (values)
+      - [PAGE_SIZE*HEAD_DIM : PAGE_BYTES) : random positive fp32 scales (bytes)
+    """
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    raw = torch.empty(num_pages, PAGE_BYTES, dtype=torch.uint8, device=device)
+
+    # Random FP8 byte pattern for the value section. Bias away from extreme
+    # bit patterns that map to NaN/Inf in float8_e4m3fn (sign|exp4|mantissa3;
+    # exp=0xF mantissa!=0 -> NaN). Restricting to [0, 0x6F] keeps |x| < 256.
+    val_bytes = torch.randint(
+        0, 0x70, (num_pages, PAGE_SIZE * HEAD_DIM), generator=g, dtype=torch.uint8
+    ).to(device)
+    raw[:, : PAGE_SIZE * HEAD_DIM] = val_bytes
+
+    # Positive fp32 scales in [0.05, 0.55]. Byte-view into the trailing region.
+    scales = (
+        torch.rand((num_pages, PAGE_SIZE), generator=g, dtype=torch.float32).to(device)
+        * 0.5
+        + 0.05
+    )
+    raw[:, PAGE_SIZE * HEAD_DIM :] = scales.contiguous().view(torch.uint8)
+
+    kv = raw.view(num_pages, PAGE_SIZE, 1, HEAD_DIM_WITH_SF)
+    return kv if dtype_view == torch.uint8 else kv.view(dtype=dtype_view)
+
+
+def _build_inputs(
+    batch_size: int,
+    seq_lens: list[int],
+    *,
+    kv_dtype_view: torch.dtype,
+    num_heads: int = 32,
+    device: torch.device = torch.device("cuda"),
+    seed: int = 0,
+):
+    assert len(seq_lens) == batch_size
+    max_seq_len = max(seq_lens)
+    max_pages = (max_seq_len + PAGE_SIZE - 1) // PAGE_SIZE
+    # One global page pool, batch picks its own page ids.
+    num_pages_total = batch_size * max_pages + 1
+
+    kvcache = _build_kvcache(
+        num_pages_total, dtype_view=kv_dtype_view, device=device, seed=seed
+    )
+
+    g = torch.Generator(device="cpu").manual_seed(seed + 1)
+    # Construct query as random bf16->fp8 to keep values inside fp8 range.
+    q_bf16 = torch.randn(
+        (batch_size, 1, num_heads, HEAD_DIM), generator=g, dtype=torch.float32
+    ).to(device)
+    q_bf16 = q_bf16.clamp_(-2.0, 2.0)
+    q_fp8 = q_bf16.to(FP8_DTYPE)
+    if kv_dtype_view == torch.uint8:
+        # Query dtype isn't toggled — only kvcache. q always fp8.
+        pass
+
+    weight = (
+        torch.rand((batch_size, num_heads), generator=g, dtype=torch.float32).to(device)
+        * 0.5
+    )
+
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+
+    # Each batch occupies its own slice of pages, randomized for realism.
+    page_table = torch.zeros((batch_size, max_pages), dtype=torch.int32, device=device)
+    for i in range(batch_size):
+        page_table[i] = torch.arange(
+            1 + i * max_pages, 1 + (i + 1) * max_pages, dtype=torch.int32, device=device
+        )
+
+    return q_fp8, kvcache, weight, seq_lens_t, page_table, max_seq_len
+
+
+def _compare(
+    ref: torch.Tensor,
+    sm120: torch.Tensor,
+    seq_lens: torch.Tensor,
+    *,
+    atol: float = 1e-3,
+    rtol: float = 1e-3,
+):
+    """Compare reference (uninitialized beyond seq_len) vs SM120 (-inf beyond)."""
+    # Valid positions must match
+    batch_size, max_seq_len = ref.shape
+    for i in range(batch_size):
+        sl = int(seq_lens[i].item())
+        torch.testing.assert_close(
+            ref[i, :sl], sm120[i, :sl], atol=atol, rtol=rtol, equal_nan=False
+        )
+    # Invalid positions in SM120 output must be -inf
+    positions = torch.arange(max_seq_len, device=sm120.device)
+    invalid = positions.unsqueeze(0) >= seq_lens.unsqueeze(1)
+    assert torch.all(
+        torch.isinf(sm120[invalid]) & (sm120[invalid] < 0)
+    ), "SM120 output must fill invalid positions with -inf"
+
+
+class TestSM120PagedMqaLogitsTorch(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA required")
+        cls.device = torch.device("cuda")
+
+    def _run_one(
+        self,
+        batch_size: int,
+        seq_lens: list[int],
+        kv_dtype_view: torch.dtype,
+        num_heads: int = 32,
+    ):
+        q, kv, w, sl, pt, msl = _build_inputs(
+            batch_size,
+            seq_lens,
+            kv_dtype_view=kv_dtype_view,
+            num_heads=num_heads,
+            device=self.device,
+        )
+        # Reference (loopy)
+        ref = fp8_paged_mqa_logits_torch(
+            q,
+            kv,
+            w,
+            sl,
+            pt,
+            deep_gemm_metadata=None,
+            max_seq_len=msl,
+            clean_logits=False,
+        )
+        # SM120 vectorized
+        sm120 = fp8_paged_mqa_logits_torch_sm120(
+            q,
+            kv,
+            w,
+            sl,
+            pt,
+            deep_gemm_metadata=None,
+            max_seq_len=msl,
+            clean_logits=False,
+        )
+        _compare(ref, sm120, sl)
+
+    def test_equiv_fp8_view_bs1(self):
+        self._run_one(1, [128], FP8_DTYPE)
+
+    def test_equiv_fp8_view_bs4_uniform(self):
+        self._run_one(4, [128, 128, 128, 128], FP8_DTYPE)
+
+    def test_equiv_fp8_view_bs4_variable(self):
+        self._run_one(4, [40, 96, 200, 256], FP8_DTYPE)
+
+    def test_equiv_uint8_view_bs1(self):
+        """Regression guard: KV cache viewed as raw uint8 (no FP8 dtype hint).
+
+        Historic bug (progress doc §4): some kernels treated uint8 bytes as raw
+        integers, producing garbled attention output. The vectorized impl must
+        produce the same numbers as the loopy reference regardless of the
+        caller's KV view dtype.
+        """
+        self._run_one(1, [128], torch.uint8)
+
+    def test_equiv_uint8_view_bs4_variable(self):
+        self._run_one(4, [40, 96, 200, 256], torch.uint8)
+
+    def test_seq_lens_zero_remainder(self):
+        """seq_len not aligned to page_size — last partial page must mask correctly."""
+        self._run_one(2, [65, 129], FP8_DTYPE)  # 65 = 1 full page + 1 token
+
+    def test_seq_lens_full_pages(self):
+        self._run_one(2, [64, 192], FP8_DTYPE)
+
+    def test_seq_lens_2d_input_accepted(self):
+        """SM120 impl squeezes seq_lens if dim>1 (matches indexer.py call site)."""
+        q, kv, w, sl, pt, msl = _build_inputs(
+            2, [64, 128], kv_dtype_view=FP8_DTYPE, device=self.device
+        )
+        sl_2d = sl.unsqueeze(-1)  # (B, 1)
+        ref = fp8_paged_mqa_logits_torch(
+            q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+        )
+        sm120 = fp8_paged_mqa_logits_torch_sm120(
+            q, kv, w, sl_2d, pt, None, max_seq_len=msl, clean_logits=False
+        )
+        _compare(ref, sm120, sl)
+
+    def test_cuda_graph_capture_and_replay(self):
+        """No .item() / data-dependent control flow — must be CUDA-graph safe."""
+        batch_size = 2
+        seq_lens = [128, 192]
+        q, kv, w, sl, pt, msl = _build_inputs(
+            batch_size, seq_lens, kv_dtype_view=FP8_DTYPE, device=self.device
+        )
+
+        # Warmup outside graph
+        for _ in range(2):
+            _ = fp8_paged_mqa_logits_torch_sm120(
+                q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+            )
+        torch.cuda.synchronize()
+
+        # Pre-allocated output (graph replay reuses this buffer)
+        static_logits_holder = {}
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            out = fp8_paged_mqa_logits_torch_sm120(
+                q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+            )
+            static_logits_holder["out"] = out
+
+        # Eager reference using the same inputs
+        ref = fp8_paged_mqa_logits_torch_sm120(
+            q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+        )
+
+        graph.replay()
+        torch.cuda.synchronize()
+        torch.testing.assert_close(
+            static_logits_holder["out"], ref, atol=1e-5, rtol=1e-5
+        )
+
+        # Replay with a different seq_lens (in-place edit of the captured tensor)
+        sl_new = torch.tensor([64, 256], dtype=torch.int32, device=self.device)
+        sl.copy_(sl_new)
+        graph.replay()
+        torch.cuda.synchronize()
+        ref_new = fp8_paged_mqa_logits_torch_sm120(
+            q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+        )
+        torch.testing.assert_close(
+            static_logits_holder["out"], ref_new, atol=1e-5, rtol=1e-5
+        )
+
+    def test_shape_assertions(self):
+        """Wrong head_dim or block_size must raise."""
+        q, kv, w, sl, pt, msl = _build_inputs(
+            1, [64], kv_dtype_view=FP8_DTYPE, device=self.device
+        )
+        # head_dim != 128
+        bad_q = q[..., :64]
+        with self.assertRaises(AssertionError):
+            fp8_paged_mqa_logits_torch_sm120(
+                bad_q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
+            )
+        # clean_logits=True not supported
+        with self.assertRaises(AssertionError):
+            fp8_paged_mqa_logits_torch_sm120(
+                q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=True
+            )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(unittest.main())