Cute-DSL FP8 MQA logits (#25220)
Co-authored-by: Mindy Li <11663212+limin2021@users.noreply.github.com> Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
co-authored by
Mindy Li
Brayden Zhong
Mindy Li
Brayden Zhong
parent
2d9f0b3317
commit
fefc1743a9
@@ -0,0 +1,361 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import sglang.jit_kernel.dsa.cutedsl_paged_mqa_logits # noqa: F401
|
||||
from sglang.jit_kernel.dsa import pick_dsl_expand
|
||||
from sglang.srt.layers.attention.dsa.utils import (
|
||||
fp8_mqa_logits_ceil_to_ue8m0,
|
||||
fp8_mqa_logits_make_fused_kv,
|
||||
)
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
|
||||
|
||||
def _generate_bench_data(
|
||||
batch_size: int,
|
||||
context_len: int,
|
||||
next_n: int,
|
||||
num_heads: int = 32,
|
||||
head_dim: int = 128,
|
||||
block_kv: int = 64,
|
||||
varlen: bool = False,
|
||||
device: str = "cuda",
|
||||
) -> dict:
|
||||
torch.manual_seed(42)
|
||||
torch.cuda.manual_seed(42)
|
||||
num_blocks_per_seq = (context_len + block_kv - 1) // block_kv
|
||||
|
||||
if varlen:
|
||||
lo = min(2048, context_len)
|
||||
context_lens = torch.randint(
|
||||
lo, context_len + 1, (batch_size,), dtype=torch.int32, device=device
|
||||
)
|
||||
total_blocks = ((context_lens + block_kv - 1) // block_kv).sum().item()
|
||||
block_table = torch.zeros(
|
||||
(batch_size, num_blocks_per_seq), dtype=torch.int32, device=device
|
||||
)
|
||||
cursor = 0
|
||||
for i in range(batch_size):
|
||||
n_blks = (context_lens[i].item() + block_kv - 1) // block_kv
|
||||
block_table[i, :n_blks] = torch.arange(
|
||||
cursor, cursor + n_blks, dtype=torch.int32, device=device
|
||||
)
|
||||
cursor += n_blks
|
||||
else:
|
||||
total_blocks = batch_size * num_blocks_per_seq
|
||||
context_lens = torch.full(
|
||||
(batch_size,), context_len, dtype=torch.int32, device=device
|
||||
)
|
||||
block_table = torch.arange(
|
||||
total_blocks, dtype=torch.int32, device=device
|
||||
).reshape(batch_size, num_blocks_per_seq)
|
||||
|
||||
q_bf16 = torch.randn(
|
||||
batch_size, next_n, num_heads, head_dim, device=device, dtype=torch.bfloat16
|
||||
)
|
||||
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
|
||||
weights = torch.randn(
|
||||
batch_size * next_n, num_heads, device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
kv_bf16 = torch.randn(
|
||||
total_blocks, block_kv, head_dim, device=device, dtype=torch.bfloat16
|
||||
)
|
||||
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
|
||||
kv_scales = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
|
||||
kv_fp8 = (kv_bf16 / kv_scales.unsqueeze(-1)).to(torch.float8_e4m3fn)
|
||||
|
||||
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scales, block_kv, head_dim)
|
||||
|
||||
return {
|
||||
"q_fp8": q_fp8,
|
||||
"kv_fp8": kv_fp8,
|
||||
"kv_scales": kv_scales,
|
||||
"kv_fused": kv_fused,
|
||||
"weights": weights,
|
||||
"context_lens": context_lens,
|
||||
"block_table": block_table,
|
||||
"max_model_len": context_len,
|
||||
"total_blocks": total_blocks,
|
||||
}
|
||||
|
||||
|
||||
def benchmark(
|
||||
batch_sizes: list[int],
|
||||
next_ns: list[int],
|
||||
context_lens: list[int],
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
varlen: bool = False,
|
||||
block_kv: int = 64,
|
||||
use_cuda_graph: bool = True,
|
||||
):
|
||||
"""Benchmark CuTe DSL FP8 paged MQA logits vs DeepGEMM.
|
||||
|
||||
The "DSL" column uses the SAME picker SGLang runs at runtime
|
||||
(``pick_dsl_expand``), so the table reflects the shipped path. For
|
||||
next_n=6 / num_heads<=42 that picker selects the native single-launch
|
||||
expansion (factor=1, atom=next_n, weights-in-SMEM) once there is enough
|
||||
work to fill the SMs, and a wave-minimizing split otherwise.
|
||||
|
||||
The "native" column force-runs the native expansion (factor=1,
|
||||
atom=next_n) whenever it fits TMEM (``next_n*num_heads <= 256``), so the
|
||||
native path is always visible even on shapes where the picker prefers a
|
||||
split. ``use_cuda_graph=True`` matches SGLang's runtime decode path.
|
||||
"""
|
||||
import deep_gemm
|
||||
from flashinfer.testing.utils import bench_gpu_time
|
||||
|
||||
num_heads = 32
|
||||
head_dim = 128
|
||||
num_sms = torch.cuda.get_device_properties(0).multi_processor_count
|
||||
|
||||
dtype_str = str(output_dtype).split(".")[-1]
|
||||
mode_str = "varlen" if varlen else "fix-len"
|
||||
backend = "CUPTI + CUDA-Graph" if use_cuda_graph else "CUPTI"
|
||||
print(
|
||||
f"output_dtype={dtype_str} mode={mode_str} block_kv={block_kv} "
|
||||
f"num_heads={num_heads} timing={backend}"
|
||||
)
|
||||
hdr = (
|
||||
f"{'batch':>5s} {'ctx':>7s} {'next_n':>6s} {'nblk':>7s} | "
|
||||
f"{'pick':>5s} {'DSL(us)':>9s} | "
|
||||
f"{'native':>6s} {'nat(us)':>9s} {'nat/DSL':>8s} | "
|
||||
f"{'DG-nat(us)':>11s} {'DG/DSL':>7s}"
|
||||
)
|
||||
print(hdr)
|
||||
print(" DSL = production picker (pick_dsl_expand) — exactly what SGLang runs")
|
||||
print(" pick = chosen factor/atom (1/6 = native single-launch; 2/3 = split)")
|
||||
print(" native = forced native expansion (factor=1, atom=next_n); '-' if N>256")
|
||||
print(" DG-nat = deep_gemm native (q=[B,next_n,H,D]) — TRT-LLM's path")
|
||||
print(" nat/DSL>1 => native beats the picker; DG/DSL>1 => DSL beats deep_gemm")
|
||||
print("-" * len(hdr))
|
||||
|
||||
# See `cutedsl_paged_mqa_logits.py` for the SPLIT_KV=256 alignment note:
|
||||
# DG metadata wrapper computes SPLIT_KV = block_kv_arg * 4 with the
|
||||
# multiplier hardcoded to 4 on SM100, and both kernels expect SPLIT_KV=256
|
||||
# (DSL: compute_tile=128 × kNumMathWarpGroups=2; DG: hardcoded). Pass 64.
|
||||
DG_METADATA_BLOCK_KV = 64
|
||||
|
||||
for next_n in next_ns:
|
||||
for context_len in context_lens:
|
||||
for batch_size in batch_sizes:
|
||||
nblk = batch_size * ((context_len + block_kv - 1) // block_kv)
|
||||
|
||||
data = _generate_bench_data(
|
||||
batch_size,
|
||||
context_len,
|
||||
next_n,
|
||||
num_heads,
|
||||
head_dim,
|
||||
block_kv,
|
||||
varlen=varlen,
|
||||
)
|
||||
|
||||
# Reshape Q + repeat ctx/block_table per (factor, atom). weights
|
||||
# [B*next_n, H] = [B*factor*atom, H] needs no reshape because the
|
||||
# row layout is preserved under [B*factor, atom, ...] view.
|
||||
def _split(factor, atom, data=data, B=batch_size):
|
||||
if factor > 1:
|
||||
return {
|
||||
"q": data["q_fp8"].reshape(
|
||||
B * factor, atom, num_heads, head_dim
|
||||
),
|
||||
"ctx_lens": data["context_lens"].repeat_interleave(factor),
|
||||
"block_table": data["block_table"].repeat_interleave(
|
||||
factor, dim=0
|
||||
),
|
||||
}
|
||||
return {
|
||||
"q": data["q_fp8"],
|
||||
"ctx_lens": data["context_lens"],
|
||||
"block_table": data["block_table"],
|
||||
}
|
||||
|
||||
def _meta(t):
|
||||
return deep_gemm.get_paged_mqa_logits_metadata(
|
||||
t["ctx_lens"].unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms
|
||||
)
|
||||
|
||||
def _make_dsl(t, schedule_meta, data=data, epi_dtype=output_dtype):
|
||||
def _dsl(
|
||||
t=t,
|
||||
schedule_meta=schedule_meta,
|
||||
data=data,
|
||||
epi_dtype=epi_dtype,
|
||||
):
|
||||
torch.ops.sglang.cute_dsl_fp8_paged_mqa_logits(
|
||||
t["q"],
|
||||
data["kv_fused"],
|
||||
data["weights"],
|
||||
t["ctx_lens"],
|
||||
t["block_table"],
|
||||
schedule_meta,
|
||||
data["max_model_len"],
|
||||
epi_dtype=epi_dtype,
|
||||
acc_dtype=output_dtype,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
|
||||
return _dsl
|
||||
|
||||
# Production pick — exactly what SGLang's runtime selects. The op
|
||||
# then auto-tunes the epilogue (incl. max_w_in_reg=8 for native).
|
||||
factor, atom = pick_dsl_expand(
|
||||
next_n, batch_size, context_len, num_sms, num_heads=num_heads
|
||||
)
|
||||
prod_t = _split(factor, atom)
|
||||
_dsl_prod = _make_dsl(prod_t, _meta(prod_t))
|
||||
|
||||
# Forced native expansion (factor=1, atom=next_n) when it fits TMEM.
|
||||
native_fits = next_n * num_heads <= 256
|
||||
if native_fits:
|
||||
nat_t = _split(1, next_n)
|
||||
_dsl_nat = _make_dsl(nat_t, _meta(nat_t))
|
||||
|
||||
# DG-native: schedule input shape (B, next_n) — wrapper
|
||||
# derives num_next_n_atoms = next_n. q stays [B, next_n, H, D].
|
||||
# On SM100, DG handles next_n ∈ {1, 2, 4} natively and falls
|
||||
# back to per-token expansion for next_n ∈ {3, 5+} internally.
|
||||
dg_nat_ctx_2d = (
|
||||
data["context_lens"].unsqueeze(-1).expand(-1, next_n).contiguous()
|
||||
)
|
||||
dg_nat_schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
|
||||
dg_nat_ctx_2d, DG_METADATA_BLOCK_KV, num_sms
|
||||
)
|
||||
|
||||
def _dg_native():
|
||||
deep_gemm.fp8_paged_mqa_logits(
|
||||
data["q_fp8"],
|
||||
data["kv_fused"],
|
||||
data["weights"],
|
||||
dg_nat_ctx_2d,
|
||||
data["block_table"],
|
||||
dg_nat_schedule_meta,
|
||||
data["max_model_len"],
|
||||
clean_logits=False,
|
||||
)
|
||||
|
||||
bench_kwargs = dict(
|
||||
dry_run_iters=5,
|
||||
repeat_iters=30,
|
||||
enable_cupti=True,
|
||||
use_cuda_graph=use_cuda_graph,
|
||||
cold_l2_cache=True,
|
||||
)
|
||||
|
||||
# First DSL call(s) trigger JIT compile; do them outside timing.
|
||||
_dsl_prod()
|
||||
if native_fits:
|
||||
_dsl_nat()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
prod_ms = np.median(bench_gpu_time(_dsl_prod, **bench_kwargs))
|
||||
nat_ms = (
|
||||
np.median(bench_gpu_time(_dsl_nat, **bench_kwargs))
|
||||
if native_fits
|
||||
else float("nan")
|
||||
)
|
||||
dg_nat_ms = np.median(bench_gpu_time(_dg_native, **bench_kwargs))
|
||||
|
||||
pick_lab = f"{factor}/{atom}"
|
||||
nat_over_prod = (
|
||||
prod_ms / nat_ms if native_fits and nat_ms > 0 else float("nan")
|
||||
)
|
||||
dg_over_prod = dg_nat_ms / prod_ms if prod_ms > 0 else float("nan")
|
||||
native_tag = f"1/{next_n}" if native_fits else "-"
|
||||
nat_us_cell = f"{nat_ms * 1e3:9.2f}" if native_fits else f"{'-':>9s}"
|
||||
nat_ratio_cell = (
|
||||
f"{nat_over_prod:7.3f}x" if native_fits else f"{'-':>8s}"
|
||||
)
|
||||
print(
|
||||
f"{batch_size:5d} {context_len:7d} {next_n:6d} {nblk:7d} | "
|
||||
f"{pick_lab:>5s} {prod_ms * 1e3:9.2f} | "
|
||||
f"{native_tag:>6s} {nat_us_cell} {nat_ratio_cell} | "
|
||||
f"{dg_nat_ms * 1e3:11.2f} {dg_over_prod:6.3f}x"
|
||||
)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark CuTe DSL FP8 paged MQA logits vs DeepGEMM "
|
||||
"(CUPTI + CUDA-graph timing via flashinfer.bench_gpu_time)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 2, 4, 6, 8, 10, 12, 14, 16],
|
||||
help="Batch sizes to sweep (default: 1 2 4 6 8 10 12 14 16).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--next_n",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 2, 4, 6],
|
||||
help="next_n values to sweep (default: 1 2 4 6 — covers decode and "
|
||||
"spec-decode target_verify with num_draft_tokens in {2, 4, 6}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context_len",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[4096, 10240, 32768, 81920, 131072],
|
||||
help="Context lengths to sweep (default: 4096 10240 32768 81920 131072).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dtype",
|
||||
type=str,
|
||||
default="float32",
|
||||
choices=["float32", "float16"],
|
||||
help="Output dtype (default: float32).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--varlen",
|
||||
action="store_true",
|
||||
help="Use variable-length per-seq context (mimics mixed-batch serving). "
|
||||
"Default is fix-length for clean comparisons.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--block_kv",
|
||||
type=int,
|
||||
default=64,
|
||||
choices=[32, 64, 128],
|
||||
help="Cache page size in tokens (default: 64 — matches SGLang NSA).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-cuda-graph",
|
||||
action="store_true",
|
||||
help="Disable CUDA-graph capture; measure launch overhead too.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not is_sm100_supported():
|
||||
print(
|
||||
"Skipping: CuTe DSL FP8 Paged MQA Logits kernel only supports "
|
||||
"SM 100 family (Blackwell)."
|
||||
)
|
||||
sys.exit(0)
|
||||
|
||||
dtype_map = {"float32": torch.float32, "float16": torch.float16}
|
||||
benchmark(
|
||||
batch_sizes=args.batch_size,
|
||||
next_ns=args.next_n,
|
||||
context_lens=args.context_len,
|
||||
output_dtype=dtype_map[args.output_dtype],
|
||||
varlen=args.varlen,
|
||||
block_kv=args.block_kv,
|
||||
use_cuda_graph=not args.no_cuda_graph,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,16 @@
|
||||
from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
|
||||
from .paged_mqa_logits import (
|
||||
aiter_paged_mqa_logits,
|
||||
cutedsl_paged_mqa_logits,
|
||||
deepgemm_paged_mqa_logits_native,
|
||||
deepgemm_paged_mqa_logits_split,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CuteDSLPagedMQALogitsRunner",
|
||||
"pick_dsl_expand",
|
||||
"aiter_paged_mqa_logits",
|
||||
"cutedsl_paged_mqa_logits",
|
||||
"deepgemm_paged_mqa_logits_native",
|
||||
"deepgemm_paged_mqa_logits_split",
|
||||
]
|
||||
@@ -0,0 +1,475 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CuTe DSL FP8 Paged MQA Logits runner and custom op.
|
||||
|
||||
Ported from TensorRT-LLM https://github.com/NVIDIA/TensorRT-LLM/pull/13219
|
||||
Provides ``torch.ops.sglang.cute_dsl_fp8_paged_mqa_logits`` as an alternative
|
||||
to ``deep_gemm.fp8_paged_mqa_logits`` on Blackwell SM100. It performs well, when the bs is low,
|
||||
and the context is long.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import torch
|
||||
from cutlass.utils import HardwareInfo
|
||||
|
||||
from sglang.jit_kernel.cutedsl_fp8_paged_mqa_logits import FP8MQALogitsKernel
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def pick_dsl_expand(
|
||||
next_n: int,
|
||||
batch_size: int = 0,
|
||||
max_ctx: int = 0,
|
||||
num_sms: int = 148,
|
||||
kernel_atoms: tuple[int, ...] = (1, 2, 3, 4),
|
||||
num_heads: int = 0,
|
||||
) -> tuple[int, int]:
|
||||
"""Pick (expand_factor, effective_next_n) for the DSL paged kernel
|
||||
using a wave-aware strategy.
|
||||
|
||||
The DSL FP8 kernel natively supports ``effective_next_n ∈ kernel_atoms``
|
||||
(default ``(1, 2, 3, 4)``). When SM utilization can be improved, reshape
|
||||
``[B, next_n, ...]`` -> ``[B * expand_factor, effective_next_n, ...]``
|
||||
caller-side.
|
||||
|
||||
Strategy: enumerate ``(expand_factor, effective_next_n)`` pairs with
|
||||
``expand_factor * effective_next_n == next_n`` and ``effective_next_n
|
||||
in kernel_atoms``. Score each by ``(waves, -expand_factor)`` where
|
||||
``waves = ceil(B * expand_factor * ceil(max_ctx/256) / num_sms)``.
|
||||
Pick min waves; on tie, prefer LARGER expand_factor (more SMs busy per
|
||||
wave; pays HBM cost of expand_factor x KV re-reads).
|
||||
|
||||
When ``batch_size == 0`` or ``max_ctx == 0`` (workload unknown), fall
|
||||
back to the legacy HBM-minimizing heuristic: largest effective_next_n
|
||||
that divides next_n cleanly (still constrained to ``kernel_atoms``).
|
||||
"""
|
||||
if batch_size <= 0 or max_ctx <= 0:
|
||||
for eff in sorted(kernel_atoms, reverse=True):
|
||||
if next_n % eff == 0:
|
||||
return next_n // eff, eff
|
||||
return next_n, 1
|
||||
|
||||
# Measured override for next_n=6, num_heads=32 on SM100 (~148 SMs): the
|
||||
# min-waves heuristic picks 3/1->2 (atom=3) but the kernel's atom=2 path
|
||||
# wins by up to ~20% in a jagged (batch, ctx) region the wave model can't
|
||||
# see. These bounds are empirical (autotuned), not analytic; outside them
|
||||
# min-waves is optimal. Reduces mean split-regret 1.0%->0.05% on the grid.
|
||||
if next_n == 6:
|
||||
# Native single-launch next_n=6 (factor=1) reads KV once vs 2-3x for the
|
||||
# split, and with weights-in-SMEM (see _propose_epi_config) it beats the
|
||||
# split by +17..44% once there is enough work to fill the SMs. Requires
|
||||
# N=6*num_heads<=256 (the single-MMA TMEM limit), i.e. num_heads<=42.
|
||||
# Below the work threshold (few SMs busy) the split's extra tasks win.
|
||||
if 0 < num_heads * 6 <= 256:
|
||||
native_wins = (
|
||||
batch_size >= 16
|
||||
or (batch_size >= 4 and max_ctx >= 32768)
|
||||
or (batch_size >= 2 and max_ctx >= 131072)
|
||||
)
|
||||
if native_wins:
|
||||
return 1, 6
|
||||
use_atom2 = (
|
||||
(batch_size >= 45 and max_ctx <= (batch_size - 44) * 32768)
|
||||
or (batch_size == 16 and max_ctx >= 49152)
|
||||
or (batch_size == 10 and max_ctx >= 90000)
|
||||
or (batch_size == 17 and 49152 <= max_ctx <= 110000)
|
||||
or (batch_size == 7 and max_ctx >= 120000)
|
||||
)
|
||||
if use_atom2 and 2 in kernel_atoms:
|
||||
return 3, 2
|
||||
|
||||
SPLIT_KV_TOKENS = 256
|
||||
cands = []
|
||||
for eff in kernel_atoms:
|
||||
if next_n % eff == 0:
|
||||
factor = next_n // eff
|
||||
ntask = (
|
||||
batch_size
|
||||
* factor
|
||||
* ((max_ctx + SPLIT_KV_TOKENS - 1) // SPLIT_KV_TOKENS)
|
||||
)
|
||||
waves = (ntask + num_sms - 1) // num_sms
|
||||
cands.append((waves, factor, eff))
|
||||
if not cands:
|
||||
return next_n, 1
|
||||
cands.sort(key=lambda x: (x[0], -x[1]))
|
||||
_, factor, eff = cands[0]
|
||||
return factor, eff
|
||||
|
||||
|
||||
_TORCH_TO_CUTLASS_DTYPE = {
|
||||
torch.float16: cutlass.Float16,
|
||||
torch.bfloat16: cutlass.BFloat16,
|
||||
torch.float32: cutlass.Float32,
|
||||
}
|
||||
|
||||
# Epilogue pipeline-flag presets for the auto-tuner (see _propose_epi_config).
|
||||
_EPI_KV_UMMA_SUB = dict(
|
||||
max_kv_pipeline=True, max_umma_pipeline=True, smem_subpartition_opt=True
|
||||
)
|
||||
_EPI_NOWAIT = dict(_EPI_KV_UMMA_SUB, remove_kv_wait_in_epilogue=True)
|
||||
|
||||
|
||||
def _propose_epi_config(
|
||||
num_heads: int,
|
||||
next_n_k: int,
|
||||
batch_split: int,
|
||||
max_ctx: int,
|
||||
num_sms: int,
|
||||
) -> tuple[int, dict]:
|
||||
"""Auto-tune (num_epi_subtiles, pipeline_flags) for the FP8 MQA epilogue.
|
||||
|
||||
Tuned on B300 (~148 SMs) for the GLM-5.2 32-head path; ``next_n_k`` and
|
||||
``batch_split`` are the post-split atom and batch reaching the kernel.
|
||||
Wins +3..14% vs the untuned base config across the next_n=6 (atom 2/3)
|
||||
grid. num_heads>32 is left at the safe baseline (no change).
|
||||
"""
|
||||
# Only the <=32-head path is tuned; leave wider indexers untouched.
|
||||
if num_heads > 32 or num_heads % 8 != 0:
|
||||
return 1, {}
|
||||
# Native next_n=6 (single launch, N=6*heads<=256): the per-token weight cache
|
||||
# is 6*heads regs and spills to local/GMEM (3x slowdown). Reading weights from
|
||||
# SMEM instead (max_w_in_reg=8) avoids the spill; combined with the 1x KV read
|
||||
# (vs the split's 2-3x) this beats the split by +17..44% at HBM-bound shapes.
|
||||
if next_n_k == 6 and num_heads * 6 <= 256:
|
||||
return 1, {"max_w_in_reg": 8}
|
||||
# num_epi_subtiles=2 interleaves LDTM with FP32 FMA on the multi-slot
|
||||
# (atom != 2) epilogue; neutral elsewhere, slightly negative on atom==2.
|
||||
nst = 2 if next_n_k != 2 else 1
|
||||
if (num_heads // nst) % 4 != 0:
|
||||
nst = 1
|
||||
waves = (batch_split * ((max_ctx + 255) // 256) + num_sms - 1) // num_sms
|
||||
# The nowait + deep KV/UMMA pipeline wins broadly below grid saturation, but
|
||||
# above ~130 waves the win flips non-monotonically with occupancy (deep_gemm
|
||||
# scheduler resonance — not modelable by waves/work/fill). Measured on B300
|
||||
# (~148 SM, hd=32, dense ctx=131k saturated sweep): the 2/3-split (atom 3)
|
||||
# resonates positively exactly when post-split batch % 24 == 0 (Bs 48,72 win
|
||||
# +10..16%; 40,44,52,56,60,64,80,88 all regress 3-6%); the 3/2-split (atom 2)
|
||||
# only resonates at Bs % 144 == 0 (144 wins +12%; 48 regresses). Whitelist
|
||||
# those; otherwise fall back to the exact baseline at saturation.
|
||||
if next_n_k == 3:
|
||||
resonant = batch_split % 24 == 0
|
||||
elif next_n_k == 2:
|
||||
resonant = batch_split % 144 == 0
|
||||
else:
|
||||
resonant = False
|
||||
if waves >= 130 and not resonant:
|
||||
return 1, {}
|
||||
# The multi-slot (atom>=3) epilogue benefits at any sub-saturation occupancy;
|
||||
# the atom==2 epilogue only benefits once there is enough work to hide the
|
||||
# flag overhead (it dominates at ~1 wave).
|
||||
flags = _EPI_NOWAIT if (next_n_k >= 3 or waves >= 50) else {}
|
||||
return nst, flags
|
||||
|
||||
|
||||
class CuteDSLPagedMQALogitsRunner:
|
||||
"""Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100).
|
||||
|
||||
Caches compiled kernels keyed by static params
|
||||
(compute_block_kv, phys_block_kv, num_heads, head_dim, next_n, num_sms).
|
||||
"""
|
||||
|
||||
kernel_cache: dict[tuple, object] = dict()
|
||||
|
||||
@classmethod
|
||||
def _compile(
|
||||
cls,
|
||||
compute_block_kv,
|
||||
phys_block_kv,
|
||||
num_heads,
|
||||
head_dim,
|
||||
next_n,
|
||||
num_sms,
|
||||
num_epi_subtiles,
|
||||
epi_dtype,
|
||||
acc_dtype,
|
||||
output_dtype,
|
||||
pipeline_flags=None,
|
||||
):
|
||||
"""Compile kernel using fake tensors + TVM FFI."""
|
||||
pipeline_flags = pipeline_flags or {}
|
||||
key = (
|
||||
compute_block_kv,
|
||||
phys_block_kv,
|
||||
num_heads,
|
||||
head_dim,
|
||||
next_n,
|
||||
num_sms,
|
||||
num_epi_subtiles,
|
||||
epi_dtype,
|
||||
acc_dtype,
|
||||
output_dtype,
|
||||
tuple(sorted(pipeline_flags.items())),
|
||||
)
|
||||
if key in cls.kernel_cache:
|
||||
return
|
||||
|
||||
to_cutlass = _TORCH_TO_CUTLASS_DTYPE
|
||||
N = next_n * num_heads
|
||||
block_bytes = phys_block_kv * (head_dim + 4)
|
||||
|
||||
sym_num_phys_blocks = cute.sym_int()
|
||||
sym_B = cute.sym_int()
|
||||
max_ctx = cute.sym_int()
|
||||
max_blocks_per_seq = cute.sym_int()
|
||||
num_ctas = cute.sym_int()
|
||||
|
||||
kv_fake = cute.runtime.make_fake_compact_tensor(
|
||||
cutlass.Uint8,
|
||||
(sym_num_phys_blocks, block_bytes),
|
||||
stride_order=(1, 0),
|
||||
)
|
||||
|
||||
q_fake = cute.runtime.make_fake_compact_tensor(
|
||||
cutlass.Uint8, (N, head_dim, sym_B), stride_order=(1, 0, 2)
|
||||
)
|
||||
|
||||
w_dtype = (
|
||||
cutlass.Float16 if epi_dtype == torch.float16 else to_cutlass[epi_dtype]
|
||||
)
|
||||
w_fake = cute.runtime.make_fake_compact_tensor(
|
||||
w_dtype, (N, sym_B), stride_order=(0, 1)
|
||||
)
|
||||
|
||||
logits_fake = cute.runtime.make_fake_tensor(
|
||||
to_cutlass[output_dtype],
|
||||
(cute.sym_int(), max_ctx),
|
||||
stride=(cute.sym_int64(), 1),
|
||||
)
|
||||
|
||||
bt_fake = cute.runtime.make_fake_compact_tensor(
|
||||
cutlass.Int32, (sym_B, max_blocks_per_seq), stride_order=(1, 0)
|
||||
)
|
||||
|
||||
cl_fake = cute.runtime.make_fake_compact_tensor(
|
||||
cutlass.Int32, (sym_B,), stride_order=(0,)
|
||||
)
|
||||
|
||||
sm_fake = cute.runtime.make_fake_compact_tensor(
|
||||
cutlass.Int32, (num_ctas, 2), stride_order=(1, 0)
|
||||
)
|
||||
|
||||
fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
|
||||
|
||||
kernel = FP8MQALogitsKernel(
|
||||
block_kv=compute_block_kv,
|
||||
phys_block_kv=phys_block_kv,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
next_n=next_n,
|
||||
num_sms=num_sms,
|
||||
num_epi_subtiles=num_epi_subtiles,
|
||||
epi_dtype=to_cutlass[epi_dtype],
|
||||
acc_dtype=to_cutlass[acc_dtype],
|
||||
output_dtype=to_cutlass[output_dtype],
|
||||
**pipeline_flags,
|
||||
)
|
||||
|
||||
compiled = cute.compile(
|
||||
kernel,
|
||||
kv_fake,
|
||||
q_fake,
|
||||
w_fake,
|
||||
logits_fake,
|
||||
bt_fake,
|
||||
cl_fake,
|
||||
sm_fake,
|
||||
cutlass.Int32(1),
|
||||
cutlass.Int32(1),
|
||||
fake_stream,
|
||||
options="--enable-tvm-ffi",
|
||||
)
|
||||
cls.kernel_cache[key] = compiled
|
||||
logger.debug(
|
||||
f"[compile cute_dsl fp8_paged_mqa_logits] {key}"
|
||||
f" kv_stages={kernel.num_kv_stages}"
|
||||
f" umma_stages={kernel.num_umma_stages}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def forward(
|
||||
cls,
|
||||
q: torch.Tensor,
|
||||
kv_fused: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
schedule_meta: torch.Tensor,
|
||||
max_context_len: int,
|
||||
num_epi_subtiles: int | None = None,
|
||||
epi_dtype: torch.dtype = torch.float32,
|
||||
acc_dtype: torch.dtype = torch.float32,
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
"""Execute FP8 paged MQA logits kernel.
|
||||
|
||||
Args:
|
||||
q: [B, next_n, H, D] FP8
|
||||
kv_fused: [num_blocks, phys_block_kv, 1, D+4] uint8
|
||||
weights: [B*next_n, H] float32
|
||||
context_lens: [B] int32
|
||||
block_table: [B, max_blocks] int32
|
||||
schedule_meta: [num_sms+1, 2] int32
|
||||
max_context_len: int
|
||||
num_epi_subtiles: epilogue sub-tile count (1, 2, or 4); None auto-tunes
|
||||
epi_dtype: epilogue compute dtype
|
||||
acc_dtype: MMA accumulator dtype
|
||||
output_dtype: output logits dtype
|
||||
Returns:
|
||||
logits: [B*next_n, max_context_len] output_dtype
|
||||
"""
|
||||
B, next_n, H, D = q.shape
|
||||
N = next_n * H
|
||||
phys_block_kv = kv_fused.shape[1]
|
||||
compute_block_kv = 128
|
||||
num_phys_blocks = kv_fused.shape[0]
|
||||
num_sms = HardwareInfo().get_device_multiprocessor_count()
|
||||
|
||||
# Auto-tune (num_epi_subtiles, pipeline flags) for the epilogue when the
|
||||
# caller leaves num_epi_subtiles unset; an explicit value disables the
|
||||
# flag auto-tuning and is honored as-is.
|
||||
auto_nst, pipeline_flags = _propose_epi_config(
|
||||
H, next_n, B, max_context_len, num_sms
|
||||
)
|
||||
if num_epi_subtiles is None:
|
||||
num_epi_subtiles = auto_nst
|
||||
else:
|
||||
pipeline_flags = {}
|
||||
|
||||
# Reshape Q: [B, next_n, H, D] -> [B, N, D] -> [N, D, B]
|
||||
q_3d = q.reshape(B, N, D).permute(1, 2, 0)
|
||||
|
||||
# Reshape weights: [B*next_n, H] -> [B, N] -> [N, B]
|
||||
if epi_dtype == torch.float16:
|
||||
# TODO: move type conversion to weight loading
|
||||
w_2d = weights.reshape(B, N).half().t()
|
||||
else:
|
||||
w_2d = weights.reshape(B, N).t()
|
||||
|
||||
# Flatten fused KV to [num_phys_blocks, block_bytes]
|
||||
kv_flat = kv_fused.reshape(num_phys_blocks, -1)
|
||||
|
||||
# Allocate output with alignment padding
|
||||
SPLIT_KV = compute_block_kv * 2 # NUM_MATH_WG = 2
|
||||
aligned_max_ctx = ((max_context_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV
|
||||
logits = torch.empty(
|
||||
(B * next_n, aligned_max_ctx),
|
||||
device=q.device,
|
||||
dtype=output_dtype,
|
||||
)
|
||||
logits = logits[:, :max_context_len]
|
||||
|
||||
key = (
|
||||
compute_block_kv,
|
||||
phys_block_kv,
|
||||
H,
|
||||
D,
|
||||
next_n,
|
||||
num_sms,
|
||||
num_epi_subtiles,
|
||||
epi_dtype,
|
||||
acc_dtype,
|
||||
output_dtype,
|
||||
tuple(sorted(pipeline_flags.items())),
|
||||
)
|
||||
if key not in cls.kernel_cache:
|
||||
cls._compile(
|
||||
compute_block_kv,
|
||||
phys_block_kv,
|
||||
H,
|
||||
D,
|
||||
next_n,
|
||||
num_sms,
|
||||
num_epi_subtiles,
|
||||
epi_dtype,
|
||||
acc_dtype,
|
||||
output_dtype,
|
||||
pipeline_flags,
|
||||
)
|
||||
compiled = cls.kernel_cache[key]
|
||||
|
||||
# FP8 q needs uint8 view to match compile-time dtype
|
||||
q_for_ffi = (
|
||||
q_3d.view(torch.uint8)
|
||||
if q_3d.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
||||
else q_3d
|
||||
)
|
||||
|
||||
compiled(
|
||||
kv_flat,
|
||||
q_for_ffi,
|
||||
w_2d,
|
||||
logits,
|
||||
block_table,
|
||||
context_lens,
|
||||
schedule_meta,
|
||||
num_phys_blocks,
|
||||
B,
|
||||
)
|
||||
return logits
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"sglang::cute_dsl_fp8_paged_mqa_logits",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def cute_dsl_fp8_paged_mqa_logits(
|
||||
q: torch.Tensor,
|
||||
kv_fused: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
schedule_meta: torch.Tensor,
|
||||
max_context_len: int,
|
||||
num_epi_subtiles: int | None = None,
|
||||
epi_dtype: torch.dtype = torch.float32,
|
||||
acc_dtype: torch.dtype = torch.float32,
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
if not is_sm100_supported():
|
||||
raise ValueError("CuteDSL FP8 Paged MQA Logits only supports SM 100 family.")
|
||||
return CuteDSLPagedMQALogitsRunner.forward(
|
||||
q,
|
||||
kv_fused,
|
||||
weights,
|
||||
context_lens,
|
||||
block_table,
|
||||
schedule_meta,
|
||||
max_context_len,
|
||||
num_epi_subtiles=num_epi_subtiles,
|
||||
epi_dtype=epi_dtype,
|
||||
acc_dtype=acc_dtype,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
|
||||
|
||||
@torch.library.register_fake("sglang::cute_dsl_fp8_paged_mqa_logits")
|
||||
def _(
|
||||
q: torch.Tensor,
|
||||
kv_fused: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
schedule_meta: torch.Tensor,
|
||||
max_context_len: int,
|
||||
num_epi_subtiles: int | None = None,
|
||||
epi_dtype: torch.dtype = torch.float32,
|
||||
acc_dtype: torch.dtype = torch.float32,
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
B = q.shape[0]
|
||||
next_n = q.shape[1]
|
||||
return torch.empty(
|
||||
B * next_n,
|
||||
max_context_len,
|
||||
dtype=output_dtype,
|
||||
device=q.device,
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def deepgemm_paged_mqa_logits_native(
|
||||
fp8_paged_mqa_logits_fn: Callable[..., torch.Tensor],
|
||||
q_fp8: torch.Tensor,
|
||||
kv_cache_fp8: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
ctx_lens_2d: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
schedule_metadata: torch.Tensor,
|
||||
max_seq_len: int,
|
||||
*,
|
||||
q_offset: int,
|
||||
B: int,
|
||||
next_n: int,
|
||||
) -> torch.Tensor:
|
||||
# block_tables[::next_n] de-expands the caller's repeat_interleave without a
|
||||
# copy (DeepGEMM only checks `stride(1) == 1`).
|
||||
return fp8_paged_mqa_logits_fn(
|
||||
q_fp8[:q_offset].view(B, next_n, q_fp8.shape[1], q_fp8.shape[2]),
|
||||
kv_cache_fp8,
|
||||
weights[:q_offset],
|
||||
ctx_lens_2d,
|
||||
block_tables[::next_n],
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
|
||||
|
||||
def deepgemm_paged_mqa_logits_split(
|
||||
fp8_paged_mqa_logits_fn: Callable[..., torch.Tensor],
|
||||
q_fp8: torch.Tensor,
|
||||
kv_cache_fp8: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
ctx_lens_2d: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
schedule_metadata: torch.Tensor,
|
||||
max_seq_len: int,
|
||||
*,
|
||||
q_offset: int,
|
||||
) -> torch.Tensor:
|
||||
q_fp8 = q_fp8.unsqueeze(1)
|
||||
return fp8_paged_mqa_logits_fn(
|
||||
q_fp8[:q_offset],
|
||||
kv_cache_fp8,
|
||||
weights[:q_offset],
|
||||
ctx_lens_2d,
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
|
||||
|
||||
def aiter_paged_mqa_logits(
|
||||
q_fp8: torch.Tensor,
|
||||
kv_cache_fp8: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
max_seq_len: int,
|
||||
*,
|
||||
preshuffle: bool,
|
||||
kv_block_size: int,
|
||||
) -> torch.Tensor:
|
||||
from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
|
||||
|
||||
q_fp8 = q_fp8.unsqueeze(1)
|
||||
batch_size, next_n, _, _ = q_fp8.shape
|
||||
logits = torch.empty(
|
||||
(batch_size * next_n, max_seq_len),
|
||||
device=q_fp8.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
deepgemm_fp8_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kv_cache_fp8,
|
||||
weights,
|
||||
logits,
|
||||
seq_lens,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
Preshuffle=preshuffle,
|
||||
KVBlockSize=kv_block_size,
|
||||
)
|
||||
return logits
|
||||
|
||||
|
||||
def cutedsl_paged_mqa_logits(
|
||||
q_fp8: torch.Tensor,
|
||||
kv_cache_fp8: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
ctx_lens_1d: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
schedule_metadata: torch.Tensor | None,
|
||||
max_seq_len: int,
|
||||
*,
|
||||
q_offset: int,
|
||||
B: int,
|
||||
next_n: int,
|
||||
is_target_verify: bool,
|
||||
dsl_expand_factor: int,
|
||||
dsl_atom: int,
|
||||
blocksize: int,
|
||||
sm_count: int,
|
||||
get_paged_mqa_logits_metadata_fn: Callable[..., torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
from sglang.jit_kernel.dsa.cutedsl_paged_mqa_logits import (
|
||||
CuteDSLPagedMQALogitsRunner,
|
||||
)
|
||||
|
||||
dsl_atom_split = dsl_expand_factor > 1 and next_n == dsl_expand_factor * dsl_atom
|
||||
if is_target_verify and dsl_atom_split:
|
||||
exp_B = B * dsl_expand_factor
|
||||
q_dsl = q_fp8[:q_offset].view(exp_B, dsl_atom, q_fp8.shape[1], q_fp8.shape[2])
|
||||
ctx_lens_1d = ctx_lens_1d.repeat_interleave(dsl_expand_factor)
|
||||
block_tables_dsl = block_tables[::next_n].repeat_interleave(
|
||||
dsl_expand_factor, dim=0
|
||||
)
|
||||
schedule_metadata = get_paged_mqa_logits_metadata_fn(
|
||||
ctx_lens_1d.unsqueeze(-1), blocksize, sm_count
|
||||
)
|
||||
elif is_target_verify and next_n >= 2:
|
||||
# Native single-launch: one task per batch entry (the kernel iterates
|
||||
# next_n internally), so the schedule must be built from B-length
|
||||
# context lens, not the caller's [B, next_n] or per-token layout.
|
||||
q_dsl = q_fp8[:q_offset].view(B, next_n, q_fp8.shape[1], q_fp8.shape[2])
|
||||
block_tables_dsl = block_tables[::next_n]
|
||||
schedule_metadata = get_paged_mqa_logits_metadata_fn(
|
||||
ctx_lens_1d.unsqueeze(-1), blocksize, sm_count
|
||||
)
|
||||
else:
|
||||
q_dsl = q_fp8[:q_offset].unsqueeze(1)
|
||||
block_tables_dsl = block_tables[:B]
|
||||
|
||||
return CuteDSLPagedMQALogitsRunner.forward(
|
||||
q_dsl,
|
||||
kv_cache_fp8.view(torch.uint8),
|
||||
weights[:q_offset],
|
||||
ctx_lens_1d,
|
||||
block_tables_dsl,
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
)
|
||||
@@ -8,12 +8,22 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.jit_kernel.dsa import (
|
||||
aiter_paged_mqa_logits,
|
||||
cutedsl_paged_mqa_logits,
|
||||
deepgemm_paged_mqa_logits_native,
|
||||
deepgemm_paged_mqa_logits_split,
|
||||
pick_dsl_expand,
|
||||
)
|
||||
from sglang.jit_kernel.fused_store_index_cache import (
|
||||
can_use_dsa_fused_store,
|
||||
fused_store_index_k_cache,
|
||||
)
|
||||
from sglang.srt.compilation.compilation_config import register_split_op
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.dsa.paged_mqa_logits_backend import (
|
||||
DSAPagedMQALogitsBackend,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsa.utils import (
|
||||
aiter_can_use_preshuffle_paged_mqa,
|
||||
is_dsa_enable_prefill_cp,
|
||||
@@ -34,7 +44,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
|
||||
get_tc_piecewise_forward_context,
|
||||
is_in_tc_piecewise_cuda_graph,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.runtime_context import get_parallel, get_server_args
|
||||
from sglang.srt.state_capturer.indexer_topk import (
|
||||
maybe_capture_indexer_topk,
|
||||
)
|
||||
@@ -431,6 +441,10 @@ class Indexer(MultiPlatformOp):
|
||||
self.scale_fmt = scale_fmt
|
||||
self.softmax_scale = self.head_dim**-0.5
|
||||
|
||||
self.paged_mqa_logits_backend = DSAPagedMQALogitsBackend.resolve(
|
||||
get_server_args().dsa_paged_mqa_logits_backend
|
||||
)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _with_real_sm_count(self):
|
||||
# When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch
|
||||
@@ -818,22 +832,35 @@ class Indexer(MultiPlatformOp):
|
||||
# Reuse pre-computed schedule metadata if available (from init_forward_metadata),
|
||||
# otherwise fall back to computing it here.
|
||||
schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None)
|
||||
|
||||
assert len(q_fp8.shape) == 3
|
||||
# attn_tp_size > 1 or MAX_LEN padding mode can leave padding in the
|
||||
# hidden states; q_offset is the real (unpadded) q length.
|
||||
q_offset = sum(metadata.get_dsa_extend_len_cpu())
|
||||
|
||||
# DG-native q=[B,next_n,H,D] is faster than expanded q=[B*next_n,1,H,D]
|
||||
# for target_verify with next_n>=2 (bigger MMA tile, fewer atoms). The
|
||||
# precomputed ctx_lens_2d's shape is the single source of truth — if
|
||||
# dsa_backend chose the per-token layout (e.g. non-SM100), fall through
|
||||
# to the expanded path.
|
||||
B = metadata.get_seqlens_int32().shape[0]
|
||||
next_n = q_offset // B if B > 0 else 0
|
||||
use_cute_dsl = (
|
||||
self.paged_mqa_logits_backend.is_cutedsl()
|
||||
and not forward_batch.forward_mode.is_draft_extend_v2()
|
||||
)
|
||||
dsl_expand_factor, dsl_atom = 1, 1
|
||||
if (
|
||||
use_cute_dsl
|
||||
and forward_batch.forward_mode.is_target_verify()
|
||||
and next_n >= 2
|
||||
):
|
||||
dsl_expand_factor, dsl_atom = pick_dsl_expand(
|
||||
next_n,
|
||||
batch_size=B,
|
||||
max_ctx=max_seq_len,
|
||||
num_sms=self.sm_count,
|
||||
kernel_atoms=(1, 2, 3, 4),
|
||||
num_heads=self.n_heads,
|
||||
)
|
||||
ctx_2d = getattr(metadata, "paged_mqa_ctx_lens_2d", None)
|
||||
use_dg_native = (
|
||||
_is_cuda
|
||||
not use_cute_dsl
|
||||
and _is_cuda
|
||||
and forward_batch.forward_mode.is_target_verify()
|
||||
and next_n >= 2
|
||||
and ctx_2d is not None
|
||||
@@ -862,51 +889,61 @@ class Indexer(MultiPlatformOp):
|
||||
assert len(weights.shape) == 3
|
||||
weights = weights.squeeze(2)
|
||||
|
||||
if _is_hip:
|
||||
from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
|
||||
|
||||
q_fp8 = q_fp8.unsqueeze(1)
|
||||
batch_size, next_n, heads, _ = q_fp8.shape
|
||||
logits = torch.empty(
|
||||
(batch_size * next_n, max_seq_len),
|
||||
device=q_fp8.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
deepgemm_fp8_paged_mqa_logits(
|
||||
if self.paged_mqa_logits_backend.is_aiter():
|
||||
logits = aiter_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kv_cache_fp8,
|
||||
weights,
|
||||
logits,
|
||||
seqlens_32,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
Preshuffle=_use_aiter_preshuffle,
|
||||
KVBlockSize=block_kv,
|
||||
preshuffle=_use_aiter_preshuffle,
|
||||
kv_block_size=block_kv,
|
||||
)
|
||||
elif use_dg_native:
|
||||
# block_tables[::next_n] de-expands dsa_backend's repeat_interleave
|
||||
# without a copy (DG only checks `stride(1) == 1`).
|
||||
logits = deep_gemm.fp8_paged_mqa_logits(
|
||||
q_fp8[:q_offset].view(B, next_n, q_fp8.shape[1], q_fp8.shape[2]),
|
||||
elif use_cute_dsl:
|
||||
logits = cutedsl_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kv_cache_fp8,
|
||||
weights[:q_offset],
|
||||
seqlens_32_2d,
|
||||
block_tables[::next_n],
|
||||
weights,
|
||||
metadata.get_seqlens_int32(),
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
clean_logits=False,
|
||||
q_offset=q_offset,
|
||||
B=B,
|
||||
next_n=next_n,
|
||||
is_target_verify=forward_batch.forward_mode.is_target_verify(),
|
||||
dsl_expand_factor=dsl_expand_factor,
|
||||
dsl_atom=dsl_atom,
|
||||
blocksize=blocksize,
|
||||
sm_count=self.sm_count,
|
||||
get_paged_mqa_logits_metadata_fn=deep_gemm.get_paged_mqa_logits_metadata,
|
||||
)
|
||||
else:
|
||||
q_fp8 = q_fp8.unsqueeze(1)
|
||||
logits = deep_gemm.fp8_paged_mqa_logits(
|
||||
q_fp8[:q_offset],
|
||||
elif use_dg_native:
|
||||
logits = deepgemm_paged_mqa_logits_native(
|
||||
deep_gemm.fp8_paged_mqa_logits,
|
||||
q_fp8,
|
||||
kv_cache_fp8,
|
||||
weights[:q_offset],
|
||||
weights,
|
||||
seqlens_32_2d,
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
clean_logits=False,
|
||||
q_offset=q_offset,
|
||||
B=B,
|
||||
next_n=next_n,
|
||||
)
|
||||
else:
|
||||
logits = deepgemm_paged_mqa_logits_split(
|
||||
deep_gemm.fp8_paged_mqa_logits,
|
||||
q_fp8,
|
||||
kv_cache_fp8,
|
||||
weights,
|
||||
seqlens_32_2d,
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_seq_len,
|
||||
q_offset=q_offset,
|
||||
)
|
||||
|
||||
# NOTE(dark): logits should be cleaned in topk_transform
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from sglang.srt.utils import is_hip, is_sm100_supported
|
||||
|
||||
|
||||
class DSAPagedMQALogitsBackend(Enum):
|
||||
DEEPGEMM = "deepgemm"
|
||||
CUTEDSL = "cutedsl"
|
||||
AITER = "aiter"
|
||||
|
||||
def is_deepgemm(self) -> bool:
|
||||
return self == DSAPagedMQALogitsBackend.DEEPGEMM
|
||||
|
||||
def is_cutedsl(self) -> bool:
|
||||
return self == DSAPagedMQALogitsBackend.CUTEDSL
|
||||
|
||||
def is_aiter(self) -> bool:
|
||||
return self == DSAPagedMQALogitsBackend.AITER
|
||||
|
||||
@staticmethod
|
||||
def resolve(value: str) -> DSAPagedMQALogitsBackend:
|
||||
if is_hip():
|
||||
if value not in ("auto", "aiter"):
|
||||
raise ValueError(
|
||||
f"dsa_paged_mqa_logits_backend={value!r} is not supported on "
|
||||
"ROCm; only 'aiter' is implemented."
|
||||
)
|
||||
return DSAPagedMQALogitsBackend.AITER
|
||||
|
||||
if value == "auto" or value == "deepgemm":
|
||||
return DSAPagedMQALogitsBackend.DEEPGEMM
|
||||
if value == "aiter":
|
||||
raise ValueError("dsa_paged_mqa_logits_backend='aiter' requires ROCm.")
|
||||
if value == "cutedsl":
|
||||
if not is_sm100_supported():
|
||||
raise ValueError(
|
||||
"dsa_paged_mqa_logits_backend='cutedsl' requires SM100 (Blackwell)."
|
||||
)
|
||||
return DSAPagedMQALogitsBackend.CUTEDSL
|
||||
raise ValueError(f"Unknown dsa_paged_mqa_logits_backend: {value!r}")
|
||||
@@ -290,3 +290,29 @@ def dsa_use_prefill_cp(forward_batch, dsa_enable_prefill_cp=None):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def fp8_mqa_logits_ceil_to_ue8m0(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
|
||||
|
||||
|
||||
def fp8_mqa_logits_make_fused_kv(
|
||||
kv_fp8: torch.Tensor,
|
||||
kv_scales: torch.Tensor,
|
||||
block_kv: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
num_phys_blocks = kv_fp8.shape[0]
|
||||
per_token_size = head_dim + 4
|
||||
block_bytes = block_kv * per_token_size
|
||||
scale_offset = block_kv * head_dim
|
||||
|
||||
fused = torch.zeros(
|
||||
num_phys_blocks, block_bytes, dtype=torch.uint8, device=kv_fp8.device
|
||||
)
|
||||
for blk in range(num_phys_blocks):
|
||||
fused[blk, :scale_offset] = kv_fp8[blk].view(torch.uint8).reshape(-1)
|
||||
fused[blk, scale_offset:] = (
|
||||
kv_scales[blk].float().contiguous().view(torch.uint8).reshape(-1)
|
||||
)
|
||||
return fused.view(num_phys_blocks, block_kv, 1, per_token_size)
|
||||
|
||||
@@ -295,6 +295,8 @@ NSA_CHOICES = DSA_CHOICES # deprecated alias
|
||||
|
||||
DSA_TOPK_BACKEND_CHOICES = ["sgl-kernel", "torch", "flashinfer"]
|
||||
|
||||
DSA_PAGED_MQA_LOGITS_BACKEND_CHOICES = ["auto", "deepgemm", "cutedsl", "aiter"]
|
||||
|
||||
MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [
|
||||
"auto",
|
||||
"no_buffer",
|
||||
@@ -1372,6 +1374,13 @@ class ServerArgs:
|
||||
resolvable=True,
|
||||
),
|
||||
] = None
|
||||
dsa_paged_mqa_logits_backend: A[
|
||||
str,
|
||||
Arg(
|
||||
help="DSA indexer paged MQA logits kernel backend. Options: 'auto' (default; DeepGEMM on CUDA, aiter on ROCm), 'deepgemm', 'cutedsl' (CuTe DSL kernel, SM 100 (Blackwell) only; wins at low batch size and long context), 'aiter' (ROCm only).",
|
||||
choices=DSA_PAGED_MQA_LOGITS_BACKEND_CHOICES,
|
||||
),
|
||||
] = "auto"
|
||||
dsa_topk_backend: A[
|
||||
str,
|
||||
Arg(
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsa import cutedsl_paged_mqa_logits, pick_dsl_expand
|
||||
from sglang.srt.layers.attention.dsa.utils import (
|
||||
fp8_mqa_logits_ceil_to_ue8m0,
|
||||
fp8_mqa_logits_make_fused_kv,
|
||||
)
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=180, suite="nightly-4-gpu-b200", nightly=True)
|
||||
|
||||
BLOCK_KV = 64
|
||||
HEAD_DIM = 128
|
||||
|
||||
|
||||
def _ref_fp8_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kv_fp8,
|
||||
kv_scales,
|
||||
weights,
|
||||
context_lens,
|
||||
block_table,
|
||||
max_model_len,
|
||||
block_kv,
|
||||
):
|
||||
B, next_n, H, D = q_fp8.shape
|
||||
device = q_fp8.device
|
||||
|
||||
logits = torch.full(
|
||||
(B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32
|
||||
)
|
||||
q_f32 = q_fp8.float()
|
||||
|
||||
for b in range(B):
|
||||
ctx_len = context_lens[b].item()
|
||||
q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device)
|
||||
w = weights[b * next_n : (b + 1) * next_n, :]
|
||||
|
||||
for blk_idx in range((ctx_len + block_kv - 1) // block_kv):
|
||||
phys_blk = block_table[b, blk_idx].item()
|
||||
k_f32 = kv_fp8[phys_blk].float()
|
||||
scales = kv_scales[phys_blk]
|
||||
|
||||
k_positions = torch.arange(
|
||||
blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device
|
||||
)
|
||||
mask = (k_positions[None, :] < ctx_len) & (
|
||||
k_positions[None, :] <= q_positions[:, None]
|
||||
)
|
||||
|
||||
qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T)
|
||||
qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device))
|
||||
qk = torch.relu(qk)
|
||||
|
||||
weighted = (w.T[:, :, None] * qk).sum(dim=0)
|
||||
weighted = weighted * scales[None, :]
|
||||
|
||||
start_pos = blk_idx * block_kv
|
||||
end_pos = start_pos + block_kv
|
||||
logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where(
|
||||
mask,
|
||||
weighted,
|
||||
torch.tensor(float("-inf"), device=device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
def _generate_test_data(
|
||||
batch_size,
|
||||
next_n,
|
||||
num_heads,
|
||||
avg_context_len,
|
||||
max_model_len,
|
||||
device="cuda",
|
||||
):
|
||||
torch.manual_seed(42)
|
||||
torch.cuda.manual_seed(42)
|
||||
context_lens = torch.randint(
|
||||
max(BLOCK_KV, int(0.7 * avg_context_len)),
|
||||
int(1.3 * avg_context_len) + 1,
|
||||
(batch_size,),
|
||||
dtype=torch.int32,
|
||||
device="cpu",
|
||||
).clamp(max=max_model_len)
|
||||
|
||||
max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
|
||||
total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item()
|
||||
num_phys_blocks = total_blocks + batch_size * 2
|
||||
|
||||
block_table = torch.full(
|
||||
(batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device
|
||||
)
|
||||
blk_offset = 0
|
||||
for i in range(batch_size):
|
||||
n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV
|
||||
block_table[i, :n_blks] = torch.arange(
|
||||
blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device
|
||||
)
|
||||
blk_offset += n_blks
|
||||
|
||||
q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device)
|
||||
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
|
||||
|
||||
kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device)
|
||||
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
|
||||
kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
|
||||
kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
|
||||
|
||||
weights = torch.randn(
|
||||
batch_size * next_n, num_heads, device=device, dtype=torch.float32
|
||||
)
|
||||
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM)
|
||||
|
||||
return {
|
||||
"q_fp8": q_fp8,
|
||||
"kv_fp8": kv_fp8,
|
||||
"kv_scales": kv_scale,
|
||||
"kv_fused": kv_fused,
|
||||
"weights": weights,
|
||||
"context_lens": context_lens.to(device),
|
||||
"block_table": block_table,
|
||||
}
|
||||
|
||||
|
||||
def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len):
|
||||
device = logits.device
|
||||
positions = torch.arange(max_model_len, device=device).unsqueeze(0)
|
||||
row_indices = torch.arange(B * next_n, device=device) // next_n
|
||||
next_n_offset = torch.arange(B * next_n, device=device) % next_n
|
||||
end_pos = context_lens[row_indices] - next_n + next_n_offset
|
||||
mask = positions <= end_pos.unsqueeze(1)
|
||||
|
||||
logits_masked = logits.float().masked_fill(~mask, 0)
|
||||
ref_masked = ref_logits.float().masked_fill(~mask, 0)
|
||||
torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def _run_cutedsl_paged_mqa_logits(
|
||||
data, batch_size, next_n, num_heads, max_model_len, is_target_verify
|
||||
):
|
||||
"""Mirrors the CUTEDSL dispatch in
|
||||
sglang.srt.layers.attention.dsa.dsa_indexer.Indexer._get_topk_paged."""
|
||||
import deep_gemm
|
||||
|
||||
num_sms = torch.cuda.get_device_properties(0).multi_processor_count
|
||||
if is_target_verify and next_n >= 2:
|
||||
dsl_expand_factor, dsl_atom = pick_dsl_expand(
|
||||
next_n,
|
||||
batch_size=batch_size,
|
||||
max_ctx=max_model_len,
|
||||
num_sms=num_sms,
|
||||
kernel_atoms=(1, 2, 3, 4),
|
||||
num_heads=num_heads,
|
||||
)
|
||||
else:
|
||||
dsl_expand_factor, dsl_atom = 1, 1
|
||||
|
||||
context_lens = data["context_lens"]
|
||||
expanded_ctx = (
|
||||
context_lens.unsqueeze(-1)
|
||||
- next_n
|
||||
+ torch.arange(1, next_n + 1, device=context_lens.device, dtype=torch.int32)
|
||||
).flatten()
|
||||
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
|
||||
expanded_ctx.unsqueeze(-1), BLOCK_KV, num_sms
|
||||
)
|
||||
block_tables_expanded = data["block_table"].repeat_interleave(next_n, dim=0)
|
||||
|
||||
return cutedsl_paged_mqa_logits(
|
||||
data["q_fp8"].view(batch_size * next_n, num_heads, HEAD_DIM),
|
||||
data["kv_fused"],
|
||||
data["weights"],
|
||||
context_lens,
|
||||
block_tables_expanded,
|
||||
schedule_metadata,
|
||||
max_model_len,
|
||||
q_offset=batch_size * next_n,
|
||||
B=batch_size,
|
||||
next_n=next_n,
|
||||
is_target_verify=is_target_verify,
|
||||
dsl_expand_factor=dsl_expand_factor,
|
||||
dsl_atom=dsl_atom,
|
||||
blocksize=BLOCK_KV,
|
||||
sm_count=num_sms,
|
||||
get_paged_mqa_logits_metadata_fn=deep_gemm.get_paged_mqa_logits_metadata,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm100_supported(),
|
||||
reason="CuTe DSL FP8 Paged MQA Logits only supports SM 100 family.",
|
||||
)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
|
||||
@pytest.mark.parametrize("next_n", [1, 2, 3, 4, 5, 6])
|
||||
@pytest.mark.parametrize("num_heads", [32, 64])
|
||||
@pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384])
|
||||
def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
|
||||
max_model_len = max(avg_ctx * 2, 2048)
|
||||
data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len)
|
||||
|
||||
logits = _run_cutedsl_paged_mqa_logits(
|
||||
data,
|
||||
batch_size,
|
||||
next_n,
|
||||
num_heads,
|
||||
max_model_len,
|
||||
is_target_verify=next_n >= 2,
|
||||
)
|
||||
|
||||
ref_logits = _ref_fp8_paged_mqa_logits(
|
||||
data["q_fp8"],
|
||||
data["kv_fp8"],
|
||||
data["kv_scales"],
|
||||
data["weights"],
|
||||
data["context_lens"],
|
||||
data["block_table"],
|
||||
max_model_len,
|
||||
BLOCK_KV,
|
||||
)
|
||||
_assert_matches_ref(
|
||||
logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,233 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsa import (
|
||||
deepgemm_paged_mqa_logits_native,
|
||||
deepgemm_paged_mqa_logits_split,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsa.utils import (
|
||||
fp8_mqa_logits_ceil_to_ue8m0,
|
||||
fp8_mqa_logits_make_fused_kv,
|
||||
)
|
||||
from sglang.srt.utils import is_sm90_supported, is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=40, suite="nightly-4-gpu-b200", nightly=True)
|
||||
|
||||
BLOCK_KV = 64
|
||||
HEAD_DIM = 128
|
||||
|
||||
|
||||
def _ref_fp8_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kv_fp8,
|
||||
kv_scales,
|
||||
weights,
|
||||
context_lens,
|
||||
block_table,
|
||||
max_model_len,
|
||||
block_kv,
|
||||
):
|
||||
B, next_n, H, D = q_fp8.shape
|
||||
device = q_fp8.device
|
||||
|
||||
logits = torch.full(
|
||||
(B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32
|
||||
)
|
||||
q_f32 = q_fp8.float()
|
||||
|
||||
for b in range(B):
|
||||
ctx_len = context_lens[b].item()
|
||||
q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device)
|
||||
w = weights[b * next_n : (b + 1) * next_n, :]
|
||||
|
||||
for blk_idx in range((ctx_len + block_kv - 1) // block_kv):
|
||||
phys_blk = block_table[b, blk_idx].item()
|
||||
k_f32 = kv_fp8[phys_blk].float()
|
||||
scales = kv_scales[phys_blk]
|
||||
|
||||
k_positions = torch.arange(
|
||||
blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device
|
||||
)
|
||||
mask = (k_positions[None, :] < ctx_len) & (
|
||||
k_positions[None, :] <= q_positions[:, None]
|
||||
)
|
||||
|
||||
qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T)
|
||||
qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device))
|
||||
qk = torch.relu(qk)
|
||||
|
||||
weighted = (w.T[:, :, None] * qk).sum(dim=0)
|
||||
weighted = weighted * scales[None, :]
|
||||
|
||||
start_pos = blk_idx * block_kv
|
||||
end_pos = start_pos + block_kv
|
||||
logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where(
|
||||
mask,
|
||||
weighted,
|
||||
torch.tensor(float("-inf"), device=device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
def _generate_test_data(
|
||||
batch_size,
|
||||
next_n,
|
||||
num_heads,
|
||||
avg_context_len,
|
||||
max_model_len,
|
||||
device="cuda",
|
||||
):
|
||||
torch.manual_seed(42)
|
||||
torch.cuda.manual_seed(42)
|
||||
context_lens = torch.randint(
|
||||
max(BLOCK_KV, int(0.7 * avg_context_len)),
|
||||
int(1.3 * avg_context_len) + 1,
|
||||
(batch_size,),
|
||||
dtype=torch.int32,
|
||||
device="cpu",
|
||||
).clamp(max=max_model_len)
|
||||
|
||||
max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
|
||||
total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item()
|
||||
num_phys_blocks = total_blocks + batch_size * 2
|
||||
|
||||
block_table = torch.full(
|
||||
(batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device
|
||||
)
|
||||
blk_offset = 0
|
||||
for i in range(batch_size):
|
||||
n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV
|
||||
block_table[i, :n_blks] = torch.arange(
|
||||
blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device
|
||||
)
|
||||
blk_offset += n_blks
|
||||
|
||||
q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device)
|
||||
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
|
||||
|
||||
kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device)
|
||||
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
|
||||
kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
|
||||
kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
|
||||
|
||||
weights = torch.randn(
|
||||
batch_size * next_n, num_heads, device=device, dtype=torch.float32
|
||||
)
|
||||
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM)
|
||||
|
||||
return {
|
||||
"q_fp8": q_fp8,
|
||||
"kv_fp8": kv_fp8,
|
||||
"kv_scales": kv_scale,
|
||||
"kv_fused": kv_fused,
|
||||
"weights": weights,
|
||||
"context_lens": context_lens.to(device),
|
||||
"block_table": block_table,
|
||||
}
|
||||
|
||||
|
||||
def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len):
|
||||
device = logits.device
|
||||
positions = torch.arange(max_model_len, device=device).unsqueeze(0)
|
||||
row_indices = torch.arange(B * next_n, device=device) // next_n
|
||||
next_n_offset = torch.arange(B * next_n, device=device) % next_n
|
||||
end_pos = context_lens[row_indices] - next_n + next_n_offset
|
||||
mask = positions <= end_pos.unsqueeze(1)
|
||||
|
||||
logits_masked = logits.float().masked_fill(~mask, 0)
|
||||
ref_masked = ref_logits.float().masked_fill(~mask, 0)
|
||||
torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def _run_deepgemm_paged_mqa_logits(data, batch_size, next_n, num_heads, max_model_len):
|
||||
"""Mirrors the DEEPGEMM dispatch in
|
||||
sglang.srt.layers.attention.dsa.dsa_indexer.Indexer._get_topk_paged:
|
||||
next_n>=2 (target-verify) goes through the native wrapper, everything
|
||||
else goes through the split wrapper."""
|
||||
import deep_gemm
|
||||
|
||||
num_sms = torch.cuda.get_device_properties(0).multi_processor_count
|
||||
|
||||
if next_n >= 2:
|
||||
ctx_lens_2d = (
|
||||
data["context_lens"].unsqueeze(-1)
|
||||
- next_n
|
||||
+ torch.arange(
|
||||
1, next_n + 1, device=data["context_lens"].device, dtype=torch.int32
|
||||
)
|
||||
)
|
||||
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
|
||||
ctx_lens_2d, BLOCK_KV, num_sms
|
||||
)
|
||||
block_tables_expanded = data["block_table"].repeat_interleave(next_n, dim=0)
|
||||
return deepgemm_paged_mqa_logits_native(
|
||||
deep_gemm.fp8_paged_mqa_logits,
|
||||
data["q_fp8"].view(batch_size * next_n, num_heads, HEAD_DIM),
|
||||
data["kv_fused"],
|
||||
data["weights"],
|
||||
ctx_lens_2d,
|
||||
block_tables_expanded,
|
||||
schedule_metadata,
|
||||
max_model_len,
|
||||
q_offset=batch_size * next_n,
|
||||
B=batch_size,
|
||||
next_n=next_n,
|
||||
)
|
||||
|
||||
ctx_lens_2d = data["context_lens"].unsqueeze(-1)
|
||||
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
|
||||
ctx_lens_2d, BLOCK_KV, num_sms
|
||||
)
|
||||
return deepgemm_paged_mqa_logits_split(
|
||||
deep_gemm.fp8_paged_mqa_logits,
|
||||
data["q_fp8"].squeeze(1),
|
||||
data["kv_fused"],
|
||||
data["weights"],
|
||||
ctx_lens_2d,
|
||||
data["block_table"],
|
||||
schedule_metadata,
|
||||
max_model_len,
|
||||
q_offset=batch_size,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm90_supported() or is_sm100_supported()),
|
||||
reason="DeepGEMM fp8_paged_mqa_logits requires SM90 (Hopper) or newer.",
|
||||
)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
|
||||
@pytest.mark.parametrize("next_n", [1, 2, 3, 4, 5, 6])
|
||||
@pytest.mark.parametrize("num_heads", [32, 64])
|
||||
@pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384])
|
||||
def test_deepgemm_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
|
||||
max_model_len = max(avg_ctx * 2, 2048)
|
||||
data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len)
|
||||
|
||||
logits = _run_deepgemm_paged_mqa_logits(
|
||||
data, batch_size, next_n, num_heads, max_model_len
|
||||
)
|
||||
|
||||
ref_logits = _ref_fp8_paged_mqa_logits(
|
||||
data["q_fp8"],
|
||||
data["kv_fp8"],
|
||||
data["kv_scales"],
|
||||
data["weights"],
|
||||
data["context_lens"],
|
||||
data["block_table"],
|
||||
max_model_len,
|
||||
BLOCK_KV,
|
||||
)
|
||||
_assert_matches_ref(
|
||||
logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -260,6 +260,7 @@ class MockModelRunner:
|
||||
"dsa_prefill_backend": "flashmla_sparse",
|
||||
"dsa_decode_backend": "fa3",
|
||||
"dsa_topk_backend": "sgl-kernel",
|
||||
"dsa_paged_mqa_logits_backend": "auto",
|
||||
},
|
||||
)()
|
||||
self.hisparse_coordinator = None
|
||||
|
||||
Reference in New Issue
Block a user