iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][SM100] Add a CuTeDSL fused query kernel (#49792)

Browse Source 
Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
xiaozhoupy
2026-08-04 15:23:06 -07:00
committed by GitHub
co-authored by Claude
parent c687c1abb8
commit 3756bf1e2b
6 changed files with 699 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_fused_q_cutedsl.py
+95
View File
@@ -0,0 +1,95 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Compare the CuTeDSL and Triton implementations of the DSA fused query op.
``fused_q`` dispatches to the CuTeDSL kernel on SM100 when the dtypes and
shapes are supported and falls back to Triton otherwise. This benchmark forces
each path so the two can be compared directly.
"""
import argparse
from unittest.mock import patch
import torch
import vllm.models.deepseek_v32.common.kernels as K
from vllm.triton_utils import triton
# GLM-5.2 / DeepSeek-V3.2 DSA dimensions.
NOPE_DIM = 512
ROPE_DIM = 64
INDEX_HEADS = 32
INDEX_HEAD_DIM = 128
MAX_POS = 16384
# num_heads is per-rank: 128 attention heads over TP 8/4/2.
NUM_HEADS = [16, 32, 64]
# MTP=5 makes a concurrency-N decode batch 6*N tokens.
NUM_TOKENS = [1, 6, 24, 96, 384, 1536, 4096]
def make_inputs(num_tokens: int, num_heads: int):
dev = "cuda"
bf16 = torch.bfloat16
return dict(
positions=torch.randint(0, MAX_POS, (num_tokens,), device=dev),
q_pe=torch.randn(num_tokens, num_heads, ROPE_DIM, device=dev, dtype=bf16),
q_pe_cos_sin_cache=torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=bf16),
index_q=torch.randn(
num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=bf16
),
index_q_cos_sin_cache=torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=bf16),
ql_nope=torch.randn(num_tokens, num_heads, NOPE_DIM, device=dev, dtype=bf16),
q_scale=torch.tensor([0.5], device=dev, dtype=torch.float32),
index_weights=torch.randn(
num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32
),
index_weights_softmax_scale=INDEX_HEAD_DIM**-0.5,
index_weights_head_scale=INDEX_HEADS**-0.5,
)
def run(kwargs, *, cutedsl: bool):
# is_fused_q_cutedsl_supported gates the dispatch inside fused_q.
from vllm.models.deepseek_v32.nvidia.ops import fused_q_cutedsl as C
with patch.object(C, "is_fused_q_cutedsl_supported", lambda *_, **__: cutedsl):
K.fused_q(**kwargs)
@triton.testing.perf_report(
[
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=NUM_TOKENS,
line_arg="provider",
line_vals=["triton", "cutedsl"],
line_names=["Triton", "CuTeDSL"],
styles=[("blue", "-"), ("green", "-")],
ylabel="us",
plot_name=f"fused_q-num_heads-{nh}",
args={"num_heads": nh},
)
for nh in NUM_HEADS
]
)
def benchmark(num_tokens, num_heads, provider):
kwargs = make_inputs(num_tokens, num_heads)
fn = lambda: run(kwargs, cutedsl=provider == "cutedsl") # noqa: E731
fn() # warm up JIT / kernel compilation before timing
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--save-path", type=str, default=None)
args = parser.parse_args()
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_cutedsl
if not (current_platform.has_device_capability(100) and has_cutedsl()):
raise SystemExit("fused_q CuTeDSL requires an SM100 device and cutlass-dsl")
benchmark.run(print_data=True, show_plots=False, save_path=args.save_path)
vllm/cute_utils/__init__.py
+1
View File
@@ -19,6 +19,7 @@ from cutlass.cutlass_dsl import T, dsl_user_op
_TORCH_TO_CUTE_DTYPE = {
torch.bfloat16: BFloat16,
torch.float8_e4m3fn: Float8E4M3FN,
torch.float32: Float32,
}
_CUTE_TO_PTX_DTYPE = {
vllm/cute_utils/cvt.py
+22 -1
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from cutlass import Constexpr, Float32, Uint32, cute
from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute
from cutlass._mlir import ir
from cutlass._mlir.dialects import llvm, vector
from cutlass.cutlass_dsl import T, dsl_user_op
@@ -38,6 +38,11 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
)
elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
# recast to 32-bit registers if needed
if data.element_type == BFloat16:
data = cute.recast_tensor(data, Uint32)
assert data.element_type == Uint32
# NOTE: the output is always 1D
size = cute.size(data.shape)
out = cute.make_rmem_tensor(size * 2, Float32)
@@ -84,6 +89,22 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
return cute.TensorSSA(vec, 2, Uint32)
@dsl_user_op
def fp32x2_to_fp8x2(a0: Float32, a1: Float32, *, loc=None, ip=None) -> Uint16:
out = llvm.inline_asm(
T.i16(),
[
a0.ir_value(loc=loc, ip=ip),
a1.ir_value(loc=loc, ip=ip),
],
"cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;",
"=h,f,f",
has_side_effects=False,
is_align_stack=False,
)
return Uint16(out)
@dsl_user_op
def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
out = llvm.inline_asm(
vllm/models/deepseek_v32/common/kernels.py
+42
View File
@@ -1,7 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
@@ -773,10 +776,12 @@ def fused_q(
``(ql_nope, q_pe)`` tuple the backend expects.
"""
assert positions.ndim == 1
assert positions.dtype == torch.int64
assert q_pe.ndim == 3
assert q_pe_cos_sin_cache.ndim == 2
assert ql_nope.ndim == 3
assert ql_nope.shape[:2] == q_pe.shape[:2]
assert q_scale.dtype == torch.float32 and q_scale.numel() == 1
num_tokens = positions.shape[0]
num_q_heads = q_pe.shape[1]
@@ -794,6 +799,23 @@ def fused_q(
assert index_weights is not None
num_index_q_heads = index_q.shape[1]
index_q_head_dim = index_q.shape[2]
# fused_q is shared with the ROCm path, and the CuTeDSL module imports
# cutlass at module scope, so only reach for it on CUDA.
cutedsl_kernel: Callable[..., None] | None = None
if current_platform.is_cuda():
from vllm.models.deepseek_v32.nvidia.ops.fused_q_cutedsl import (
fused_q_cutedsl,
is_fused_q_cutedsl_supported,
)
if is_fused_q_cutedsl_supported(
q_pe,
index_q,
ql_nope,
has_indexer=has_indexer,
quantize_mqa=quantize_mqa,
):
cutedsl_kernel = fused_q_cutedsl
grid_heads = max(mqa_grid_heads, num_index_q_heads)
if quantize_mqa:
# fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
@@ -815,6 +837,26 @@ def fused_q(
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
if cutedsl_kernel is not None:
cutedsl_kernel(
positions,
q_pe,
q_pe_cos_sin_cache,
ql_nope,
q_scale,
mqa_q,
index_q,
index_q_cos_sin_cache,
index_weights,
index_weights_softmax_scale,
index_weights_head_scale,
index_q_fp8,
index_weights_out,
has_indexer=has_indexer,
index_rope_interleave=index_rope_interleave,
)
return index_q_fp8, index_weights_out, mqa_q
_fused_q_kernel[(3, num_tokens, grid_heads)](
positions,
q_pe,
vllm/models/deepseek_v32/nvidia/ops/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py
+537
View File
@@ -0,0 +1,537 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from functools import cache
import cutlass
import cutlass.cute as cute
import torch
from cuda.bindings.driver import CUstream
from cutlass import BFloat16, Float8E4M3FN, Float32, Int64, Uint8, Uint16, Uint32
from vllm.cute_utils import _TORCH_TO_CUTE_DTYPE, cvt
from vllm.platforms import current_platform
def _make_fake_tensor(dtype, shape, divisibility):
stride = tuple(
1 if i == len(shape) - 1 else cute.sym_int64(divisibility=divisibility)
for i in range(len(shape))
)
return cute.runtime.make_fake_tensor(
dtype,
shape,
stride,
assumed_align=divisibility * dtype.width // 8,
)
def is_fused_q_cutedsl_supported(
q_pe: torch.Tensor,
index_q: torch.Tensor | None,
ql_nope: torch.Tensor,
*,
has_indexer: bool,
quantize_mqa: bool,
) -> bool:
if not (
current_platform.has_device_capability(100)
and quantize_mqa
and q_pe.dtype == ql_nope.dtype == torch.bfloat16
and q_pe.shape[-1] == 64
and ql_nope.shape[-1] == 512
# One warp per head group; a non-multiple would trip the kernel's own
# assert instead of falling back to Triton.
and q_pe.shape[1] % 4 == 0
):
return False
return not has_indexer or (
index_q is not None
and index_q.dtype == torch.bfloat16
and index_q.shape[1] % 16 == 0
and index_q.shape[-1] == 128
)
def fused_q_cutedsl(
positions: torch.Tensor,
q_pe: torch.Tensor,
rope_cache: torch.Tensor,
ql_nope: torch.Tensor,
q_scale: torch.Tensor,
mqa_output: torch.Tensor,
idx_q: torch.Tensor,
idx_rope_cache: torch.Tensor,
idx_weights: torch.Tensor,
idx_weights_softmax_scale: float,
idx_weights_head_scale: float,
idx_q_fp8: torch.Tensor,
idx_weights_out: torch.Tensor,
has_indexer: bool = True,
index_rope_interleave: bool = True,
) -> None:
_, num_heads, rope_dim = q_pe.shape
_, _, nope_dim = ql_nope.shape
_, num_idx_heads, idx_dim = idx_q.shape
if has_indexer:
idx_rope_type = _TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype]
idx_weights_type = _TORCH_TO_CUTE_DTYPE[idx_weights.dtype]
else:
idx_dim = num_idx_heads = 0
idx_q = idx_rope_cache = idx_q_fp8 = None
idx_weights = idx_weights_out = None
idx_rope_type = idx_weights_type = None
rope_type = _TORCH_TO_CUTE_DTYPE[rope_cache.dtype]
compiled = FusedQKernel.compile(
rope_dim,
nope_dim,
num_heads,
rope_type,
idx_dim,
num_idx_heads,
idx_rope_type,
idx_weights_type,
index_rope_interleave,
)
compiled(
positions,
q_pe,
rope_cache,
ql_nope,
q_scale.view(1),
mqa_output,
idx_q,
idx_rope_cache,
idx_weights,
idx_q_fp8,
idx_weights_out,
float(idx_weights_softmax_scale * idx_weights_head_scale),
)
class FusedQKernel:
def __init__(
self,
rope_dim: int,
nope_dim: int,
num_heads: int,
idx_dim: int,
num_idx_heads: int,
index_rope_interleave: bool,
) -> None:
assert rope_dim == 64
assert nope_dim == 512
assert idx_dim in (128, 0)
self.rope_dim = rope_dim
self.nope_dim = nope_dim
self.num_heads = num_heads
self.idx_dim = idx_dim
self.num_idx_heads = num_idx_heads
self.index_rope_interleave = index_rope_interleave
# mqa: rope_dim=64, nope_dim=512, num_heads=64/TP
# indexer: rope_dim=64, nope_dim=64, num_heads=32
self.num_warps = 4
assert num_heads % self.num_warps == 0
assert num_idx_heads % (4 * self.num_warps) == 0
self.num_ctas_per_tok = num_heads // self.num_warps
self.num_ctas_per_idx_tok = num_idx_heads // (4 * self.num_warps)
@cute.jit
def __call__(
self,
positions: cute.Tensor,
q_pe: cute.Tensor,
rope_cache: cute.Tensor,
ql_nope: cute.Tensor,
q_scale: cute.Tensor,
mqa_output: cute.Tensor,
idx_q: cute.Tensor,
idx_rope_cache: cute.Tensor,
idx_weights: cute.Tensor,
idx_q_fp8: cute.Tensor,
idx_weights_out: cute.Tensor,
weight_scale: Float32,
stream: CUstream,
):
num_tokens = positions.shape[0]
if cutlass.const_expr(self.idx_dim == 0):
grid = (num_tokens, self.num_ctas_per_tok, 1)
else:
num_mqa_ctas = num_tokens * self.num_ctas_per_tok
num_idx_ctas = num_tokens * self.num_ctas_per_idx_tok
grid = (num_mqa_ctas + num_idx_ctas, 1, 1)
self.kernel(
positions,
q_pe,
rope_cache,
ql_nope,
q_scale,
mqa_output,
idx_q,
idx_rope_cache,
idx_weights,
idx_q_fp8,
idx_weights_out,
weight_scale,
).launch(
grid=grid,
block=(self.num_warps * 32, 1, 1),
stream=stream,
use_pdl=True,
)
@cute.kernel
def kernel(
self,
positions: cute.Tensor,
q_pe: cute.Tensor,
q_pe_rope_cache: cute.Tensor,
ql_nope: cute.Tensor,
q_scale: cute.Tensor,
mqa_output: cute.Tensor,
idx_q: cute.Tensor,
idx_q_rope_cache: cute.Tensor,
idx_weights: cute.Tensor,
idx_q_fp8: cute.Tensor,
idx_weights_out: cute.Tensor,
weight_scale: Float32,
):
if cutlass.const_expr(self.idx_dim == 0):
token_id, group_id, _ = cute.arch.block_idx()
self.mqa(
positions,
q_pe,
q_pe_rope_cache,
ql_nope,
q_scale,
mqa_output,
token_id,
group_id,
)
else:
# CTA-specialization
bid, _, _ = cute.arch.block_idx()
num_mqa_ctas = positions.shape[0] * self.num_ctas_per_tok
if bid < num_mqa_ctas:
self.mqa(
positions,
q_pe,
q_pe_rope_cache,
ql_nope,
q_scale,
mqa_output,
bid // self.num_ctas_per_tok,
bid % self.num_ctas_per_tok,
)
else:
bid -= num_mqa_ctas
self.indexer(
positions,
idx_q,
idx_q_rope_cache,
idx_weights,
idx_q_fp8,
idx_weights_out,
weight_scale,
bid // self.num_ctas_per_idx_tok,
bid % self.num_ctas_per_idx_tok,
)
@cute.jit
def mqa(
self,
positions: cute.Tensor,
q_pe: cute.Tensor,
q_pe_rope_cache: cute.Tensor,
ql_nope: cute.Tensor,
q_scale: cute.Tensor,
mqa_output: cute.Tensor,
token_id,
group_id,
):
tid, _, _ = cute.arch.thread_idx()
warp_id = cute.arch.make_warp_uniform(tid // 32)
lane_id = tid % 32
head_id = group_id * self.num_warps + warp_id
cute.arch.griddepcontrol_wait()
pos = positions[token_id]
inv_scale = 1.0 / q_scale[0]
cp_op = cute.nvgpu.CopyUniversalOp()
cp_32B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
cp_2B = cute.make_copy_atom(cp_op, Uint8, num_bits_per_copy=16)
##### issue all loads asap #####
rQ_nope_bf16 = cute.make_rmem_tensor(16, BFloat16)
rQ_rope_bf16 = cute.make_rmem_tensor(2, BFloat16)
src_ql_nope = cute.local_tile(
ql_nope[token_id, head_id, None], (16,), (lane_id,)
)
src_q_rope = cute.local_tile(q_pe[token_id, head_id, None], (2,), (lane_id,))
cute.copy(cp_32B, src_ql_nope, rQ_nope_bf16)
cute.copy(cp_4B, src_q_rope, rQ_rope_bf16)
rCos_raw = q_pe_rope_cache[pos, 0 + lane_id]
rSin_raw = q_pe_rope_cache[pos, 32 + lane_id]
##### process NoPE #####
rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16).load() * inv_scale
rQ_nope_f8 = cute.make_rmem_tensor(16, Float8E4M3FN)
rQ_nope_f8.store(rQ_nope_f32.to(Float8E4M3FN))
dst_Q_nope = cute.local_tile(
mqa_output[token_id, head_id, None], (16,), (lane_id,)
)
cute.copy(cp_16B, rQ_nope_f8, dst_Q_nope)
##### process RoPE ######
rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
rCos = rCos_raw.to(Float32)
rSin = rSin_raw.to(Float32)
r0 = (rQ_rope_f32[0] * rCos - rQ_rope_f32[1] * rSin) * inv_scale
r1 = (rQ_rope_f32[1] * rCos + rQ_rope_f32[0] * rSin) * inv_scale
cute.arch.griddepcontrol_launch_dependents()
# TensorSSA fp32->fp8 cvt has a bug. rely on direct PTX
rQ_rope_f8 = cute.make_rmem_tensor(2, Float8E4M3FN)
cute.recast_tensor(rQ_rope_f8, Uint16)[0] = cvt.fp32x2_to_fp8x2(r0, r1)
dst_Q_rope = cute.local_tile(
mqa_output[token_id, head_id, None], (2,), (256 + lane_id,)
)
cute.copy(cp_2B, rQ_rope_f8, dst_Q_rope)
@cute.jit
def indexer(
self,
positions: cute.Tensor,
idx_q: cute.Tensor,
idx_q_rope_cache: cute.Tensor,
idx_weights: cute.Tensor,
idx_q_fp8: cute.Tensor,
idx_weights_out: cute.Tensor,
weight_scale: Float32,
token_id,
group_id,
):
tid, _, _ = cute.arch.thread_idx()
subwarp_id = tid // 8
sublane_id = tid % 8
head_id = group_id * (4 * self.num_warps) + subwarp_id
cute.arch.griddepcontrol_wait()
pos = positions[token_id]
cp_op = cute.nvgpu.CopyUniversalOp()
cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
cp_8B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
##### issue all loads first #####
rQ_rope_bf16 = cute.make_rmem_tensor(8, BFloat16)
if cutlass.const_expr(self.index_rope_interleave):
src_idx_q_rope = cute.local_tile(
idx_q[token_id, head_id, None], (8,), (sublane_id,)
)
cute.copy(cp_16B, src_idx_q_rope, rQ_rope_bf16)
else:
src_idx_q_rope = cute.zipped_divide(
idx_q[token_id, head_id, None], (4,)
) # (4,32)
cute.copy(
cp_8B,
src_idx_q_rope[None, 0 + sublane_id],
cute.local_tile(rQ_rope_bf16, (4,), (0,)),
)
cute.copy(
cp_8B,
src_idx_q_rope[None, 8 + sublane_id],
cute.local_tile(rQ_rope_bf16, (4,), (1,)),
)
rQ_nope_bf16 = cute.make_rmem_tensor(8, BFloat16)
src_idx_q_nope = cute.local_tile(
idx_q[token_id, head_id, None], (8,), (8 + sublane_id,)
)
cute.copy(cp_16B, src_idx_q_nope, rQ_nope_bf16)
rCos_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
rSin_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
rope_cache_view = cute.zipped_divide(idx_q_rope_cache[pos, None], (4,))
if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
cute.copy(cp_16B, rope_cache_view[None, sublane_id], rCos_raw)
cute.copy(cp_16B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
cute.copy(cp_8B, rope_cache_view[None, sublane_id], rCos_raw)
cute.copy(cp_8B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
# unpack to FP32
rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16)
if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
rCos = rCos_raw
rSin = rSin_raw
elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
rCos = cvt.bf16x2_to_fp32x2(rCos_raw)
rSin = cvt.bf16x2_to_fp32x2(rSin_raw)
# apply rope
for i in cutlass.range_constexpr(4):
if cutlass.const_expr(self.index_rope_interleave):
r0 = rQ_rope_f32[i * 2 + 0] * rCos[i] - rQ_rope_f32[i * 2 + 1] * rSin[i]
r1 = rQ_rope_f32[i * 2 + 1] * rCos[i] + rQ_rope_f32[i * 2 + 0] * rSin[i]
rQ_rope_f32[i * 2 + 0] = r0
rQ_rope_f32[i * 2 + 1] = r1
else:
r0 = rQ_rope_f32[0 + i] * rCos[i] - rQ_rope_f32[4 + i] * rSin[i]
r1 = rQ_rope_f32[4 + i] * rCos[i] + rQ_rope_f32[0 + i] * rSin[i]
rQ_rope_f32[0 + i] = r0
rQ_rope_f32[4 + i] = r1
# amax
amax = Float32(1e-4)
for i in cutlass.range_constexpr(8):
amax = cute.arch.fmax(amax, cute.math.absf(rQ_rope_f32[i]))
amax = cute.arch.fmax(amax, cute.math.absf(rQ_nope_f32[i]))
# warp reduction among 8 lanes
for i in cutlass.range_constexpr(3):
other = cute.arch.shuffle_sync_bfly(amax, 1 << i)
amax = cute.arch.fmax(amax, other)
# compute scale from amax
# exp2(ceil(log2(scale))) via bit manipulation
scale = amax * (1.0 / 448.0)
bits = scale.bitcast(Uint32)
exp_bits = (bits + Uint32(0x007FFFFF)) & Uint32(0x7F800000)
scale = exp_bits.bitcast(Float32)
inv_scale = (Uint32(0x7F000000) - exp_bits).bitcast(Float32)
for i in cutlass.range_constexpr(8):
rQ_nope_f32[i] *= inv_scale
rQ_rope_f32[i] *= inv_scale
cute.arch.griddepcontrol_launch_dependents()
# quantize and store
rQ_nope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
rQ_nope_f8.store(rQ_nope_f32.load().to(Float8E4M3FN))
dst_idx_q_nope = cute.local_tile(
idx_q_fp8[token_id, head_id, None], (8,), (8 + sublane_id,)
)
cute.copy(cp_8B, rQ_nope_f8, dst_idx_q_nope)
rQ_rope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
rQ_rope_f8.store(rQ_rope_f32.load().to(Float8E4M3FN))
if cutlass.const_expr(self.index_rope_interleave):
dst_idx_q_rope = cute.local_tile(
idx_q_fp8[token_id, head_id, None], (8,), (sublane_id,)
)
cute.copy(cp_8B, rQ_rope_f8, dst_idx_q_rope)
else:
dst_idx_q_rope = cute.zipped_divide(
idx_q_fp8[token_id, head_id, None], (4,)
) # (4,32)
cute.copy(
cp_4B,
cute.local_tile(rQ_rope_f8, (4,), (0,)),
dst_idx_q_rope[None, 0 + sublane_id],
)
cute.copy(
cp_4B,
cute.local_tile(rQ_rope_f8, (4,), (1,)),
dst_idx_q_rope[None, 8 + sublane_id],
)
# scale indexer weights
if sublane_id == 0:
w = idx_weights[token_id, head_id].to(Float32)
idx_weights_out[token_id, head_id] = w * scale * weight_scale
@staticmethod
@cache
def compile(
rope_dim: int,
nope_dim: int,
num_heads: int,
rope_type: type[cutlass.Numeric],
idx_dim: int,
num_idx_heads: int,
idx_rope_type: type[cutlass.Numeric] | None,
idx_weights_type: type[cutlass.Numeric] | None,
index_rope_interleave: bool,
):
num_tokens = cute.sym_int()
max_pos = cute.sym_int()
positions = _make_fake_tensor(Int64, (num_tokens,), divisibility=1)
q_pe = _make_fake_tensor(
BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16
)
rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8)
ql_nope = _make_fake_tensor(
BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16
)
q_scale = _make_fake_tensor(Float32, (1,), divisibility=4)
mqa_output = _make_fake_tensor(
Float8E4M3FN,
(num_tokens, num_heads, nope_dim + rope_dim),
divisibility=16,
)
if idx_rope_type is not None:
index_q = _make_fake_tensor(
BFloat16, (num_tokens, num_idx_heads, idx_dim), divisibility=16
)
index_rope_cache = _make_fake_tensor(
idx_rope_type, (max_pos, rope_dim), divisibility=8
)
index_weights = _make_fake_tensor(
idx_weights_type, (num_tokens, num_idx_heads), divisibility=8
)
index_q_fp8 = _make_fake_tensor(
Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16
)
index_weights_out = _make_fake_tensor(
Float32, (num_tokens, num_idx_heads), divisibility=4
)
else:
index_q = index_rope_cache = index_q_fp8 = None
index_weights = index_weights_out = None
kernel = FusedQKernel(
rope_dim,
nope_dim,
num_heads,
idx_dim,
num_idx_heads,
index_rope_interleave,
)
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
return cute.compile(
kernel,
positions,
q_pe,
rope_cache,
ql_nope,
q_scale,
mqa_output,
index_q,
index_rope_cache,
index_weights,
index_q_fp8,
index_weights_out,
Float32(0.0),
stream,
options="--enable-tvm-ffi",
)