minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
eeecho
2026-06-01 14:05:20 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent dfa1af99f5
commit 524ba10eda
11 changed files with 2104 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
+47 -1
View File
@@ -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
+40 -17
View File
@@ -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(
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 = flash_mla_with_kvcache_sm120(
q=q,
k_cache=swa_k_cache,
head_dim_v=self.head_dim_v,
softmax_scale=self.softmax_scale,
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]
else:
import sgl_kernel.flash_mla as flash_mla
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
+73 -1
View File
@@ -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
+252
View File
@@ -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
+370
View File
@@ -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
+3
View File
@@ -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
+454
View File
@@ -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
+72 -2
View File
@@ -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
+14
View File
@@ -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
+465
View File
@@ -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
+314
View File
@@ -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())