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

[DSA] Use cos_sin_cache for DSA indexer fusion (#29613)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-06-29 23:49:00 -07:00
committed by GitHub
parent dbb51c46ac
commit b6fceaa789
6 changed files with 160 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh
+27 -13
View File
@@ -1,7 +1,7 @@
// DeepSeek-V3.2 only.
//
// DSA indexer K kernels: single-head LayerNorm (not RMS), ropes the leading
// kRopeDim dims, and fp8-quantizes the un-rotated activations. V3.2 drops the
// kRopeDim dims, and fp8-quantizes the rotated activations. V3.2 drops the
// Hadamard incoherence rotation; it is logit-preserving (see main_norm_rope.cuh).
//
// Independent of the wk + weights_proj GEMM fusion (dsa_indexer.py): `k_input`
@@ -33,13 +33,27 @@ constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kW
#define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// Indexer K: LayerNorm + RoPE -> bf16.
struct FusedKIndexerNormRopeParams {
const void* __restrict__ k_input; // (B, 128) DType
void* __restrict__ k_out; // (B, 128) DType
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` in elements (caller passes the wk slice directly).
int64_t k_input_stride_batch;
@@ -71,7 +85,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
@@ -82,7 +96,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
@@ -144,7 +158,7 @@ struct FusedKIndexerNormRopeKernel {
const tvm::ffi::TensorView k_out,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
@@ -176,7 +190,7 @@ struct FusedKIndexerNormRopeKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -191,7 +205,7 @@ struct FusedKIndexerNormRopeKernel {
.k_out = k_out.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
@@ -215,7 +229,7 @@ struct FusedKIndexerNormRopeStoreParams {
const void* __restrict__ indices; // (B,) int64 -- out_cache_loc
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` (caller passes the non-contiguous wk slice directly).
int64_t k_input_stride_batch;
@@ -249,7 +263,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
@@ -260,7 +274,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
@@ -345,7 +359,7 @@ struct FusedKIndexerNormRopeStoreKernel {
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
@@ -381,7 +395,7 @@ struct FusedKIndexerNormRopeStoreKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -397,7 +411,7 @@ struct FusedKIndexerNormRopeStoreKernel {
.indices = indices.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh
+29 -6
View File
@@ -46,6 +46,20 @@ constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads;
#define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16)
#define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// ============================================================================
// Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
// ============================================================================
@@ -423,7 +437,10 @@ struct FusedQIndexerRopeHadamardQuantParams {
const void* __restrict__ weight; // (B, num_heads) DType
float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat)
float weight_scale; // scalar c4_indexer.weight_scale
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
// Template-dependent layout:
// kRopeFirst=false: (max_pos, 64) fp32 interleaved [cos0, sin0, ...]
// kRopeFirst=true : (max_pos, 64) fp32 halves [cos..., sin...]
const float* __restrict__ rope_cache;
const void* __restrict__ positions; // (B,) PosT
// Row stride for `weight` (caller passes the non-contiguous wk slice directly).
int64_t weight_stride_batch;
@@ -460,7 +477,7 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
const uint32_t batch_id = work_id / params.num_heads;
const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
const auto rope_cache = params.rope_cache + position * kRopeDim;
// Lane 0 prefetches the weight scalar for this (token, head) work item.
// Weight is (B, num_heads) DType; we need one scalar per warp -- offload
@@ -477,7 +494,13 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
if (is_rope_lane) freq.load(freqs_cis, kRopeFirst ? lane_id : (lane_id - (kWarpThreads - kRopeSize)));
if (is_rope_lane) {
if constexpr (kRopeFirst) {
freq = load_rope_first_cos_sin<kRopeDim>(rope_cache, lane_id);
} else {
freq.load(rope_cache, lane_id - (kWarpThreads - kRopeSize));
}
}
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
@@ -567,7 +590,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView weights_out,
double weight_scale,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView rope_cache,
const tvm::ffi::TensorView positions) {
using namespace host;
constexpr int64_t kHeadDim = 128;
@@ -603,7 +626,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(rope_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -632,7 +655,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
.weight = weight.data_ptr(),
.weights_out = static_cast<float*>(weights_out.data_ptr()),
.weight_scale = static_cast<float>(weight_scale),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.rope_cache = static_cast<const float*>(rope_cache.data_ptr()),
.positions = positions.data_ptr(),
.weight_stride_batch = weight.stride(0),
.batch_size = batch_size,
python/sglang/jit_kernel/dsv32/elementwise.py
+4 -6
View File
@@ -43,11 +43,10 @@ def fused_k_indexer_norm_rope(
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
) -> torch.Tensor:
"""V3.2 indexer K: LayerNorm + RoPE on leading dims -> bf16. CUDA only."""
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
# k_input may be a non-contiguous wk slice; output is always contiguous.
k_out = torch.empty(k_input.shape, dtype=k_input.dtype, device=k_input.device)
module = _jit_k_indexer_norm_rope_module(k_input.dtype)
@@ -56,7 +55,7 @@ def fused_k_indexer_norm_rope(
k_out,
weight,
bias,
freqs_real,
cos_sin_cache,
positions,
float(eps),
)
@@ -70,13 +69,12 @@ def fused_k_indexer_norm_rope_store(
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
page_size: int,
) -> None:
"""V3.2 indexer K + fused store: LayerNorm + RoPE on leading dims + fp8
act-quant + paged index-k cache write, in one launch. CUDA only."""
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
if not out_cache_loc.is_contiguous():
out_cache_loc = out_cache_loc.contiguous()
module = _jit_k_indexer_norm_rope_store_module(k_input.dtype, page_size)
@@ -86,7 +84,7 @@ def fused_k_indexer_norm_rope_store(
out_cache_loc,
weight,
bias,
freqs_real,
cos_sin_cache,
positions,
float(eps),
)
python/sglang/jit_kernel/dsv4/elementwise.py
+2 -3
View File
@@ -188,11 +188,10 @@ def fused_q_indexer_rope_first_quant(
q_input: torch.Tensor,
weight: torch.Tensor,
weight_scale: float,
freqs_cis: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""DeepSeek-V3.2 only. Indexer Q: RoPE on the leading dims + fp8 act-quant. CUDA only."""
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device)
weights_out = torch.empty(
(*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device
@@ -204,7 +203,7 @@ def fused_q_indexer_rope_first_quant(
weight,
weights_out,
float(weight_scale),
freqs_real,
cos_sin_cache,
positions,
)
return q_fp8, weights_out
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+8 -18
View File
@@ -327,16 +327,6 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
return hadamard_transform(x, scale=hidden_size**-0.5)
def _shared_indexer_freqs_cis(rotary_emb: torch.nn.Module) -> torch.Tensor:
cached = getattr(rotary_emb, "_dsa_indexer_freqs_cis", None)
if cached is None:
c = rotary_emb.cos_sin_cache.to(torch.float32)
half = c.shape[-1] // 2
cached = torch.complex(c[:, :half].contiguous(), c[:, half:].contiguous())
rotary_emb._dsa_indexer_freqs_cis = cached
return cached
class Indexer(MultiPlatformOp):
_MQA_LOGITS_BYTES_PER_ELEM = 4
_MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
@@ -437,10 +427,6 @@ class Indexer(MultiPlatformOp):
self.scale_fmt = scale_fmt
self.softmax_scale = self.head_dim**-0.5
self._indexer_freqs_cis: Optional[torch.Tensor] = None
if _use_dsa_indexer_fusion:
self._indexer_freqs_cis = _shared_indexer_freqs_cis(self.rotary_emb)
@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
@@ -456,6 +442,10 @@ class Indexer(MultiPlatformOp):
else:
yield
@property
def _indexer_cos_sin_cache(self) -> torch.Tensor:
return self.rotary_emb.cos_sin_cache
def _weights_proj_bf16_in_fp32_out(
self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]]
) -> torch.Tensor:
@@ -662,7 +652,7 @@ class Indexer(MultiPlatformOp):
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
page_size,
)
@@ -674,7 +664,7 @@ class Indexer(MultiPlatformOp):
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)
self._store_index_k_cache(
@@ -726,7 +716,7 @@ class Indexer(MultiPlatformOp):
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)
@@ -763,7 +753,7 @@ class Indexer(MultiPlatformOp):
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)
test/registered/jit/test_dsv32_indexer_fusion.py
+75 -12
View File
@@ -56,9 +56,9 @@ def _make_inputs(B, seed=0, pos_dtype=torch.int32):
dev = "cuda"
cos = torch.randn(MAX_POS, HALF, device=dev, generator=g)
sin = torch.randn(MAX_POS, HALF, device=dev, generator=g)
freqs_cis = torch.complex(cos, sin)
cos_sin_cache = torch.cat((cos, sin), dim=-1)
positions = torch.randint(0, 4096, (B,), device=dev, dtype=pos_dtype, generator=g)
return cos, sin, freqs_cis, positions
return cos, sin, cos_sin_cache, positions
def _rope_first(x, cos_p, sin_p):
@@ -78,12 +78,12 @@ def test_k_norm_rope_matches_reference():
_skip_if_unavailable()
dev = "cuda"
B = 37
cos, sin, freqs_cis, positions = _make_inputs(B)
cos, sin, cos_sin_cache, positions = _make_inputs(B)
key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev)
weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
out = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
out = fused_k_indexer_norm_rope(key, weight, bias, EPS, cos_sin_cache, positions)
torch.cuda.synchronize()
normed = torch.nn.functional.layer_norm(
@@ -106,7 +106,7 @@ def test_k_store_matches_unfused(strided):
pytest.skip("fused store JIT unavailable")
dev = "cuda"
B, n_heads = 41, 64
cos, sin, freqs_cis, positions = _make_inputs(B)
cos, sin, cos_sin_cache, positions = _make_inputs(B)
weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
@@ -125,11 +125,13 @@ def test_k_store_matches_unfused(strided):
)
buf_fused = torch.zeros_like(buf_ref)
key_bf16 = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
key_bf16 = fused_k_indexer_norm_rope(
key, weight, bias, EPS, cos_sin_cache, positions
)
fused_store_index_k_cache(key_bf16, buf_ref, loc, PAGE_SIZE)
fused_k_indexer_norm_rope_store(
key, buf_fused, loc, weight, bias, EPS, freqs_cis, positions, PAGE_SIZE
key, buf_fused, loc, weight, bias, EPS, cos_sin_cache, positions, PAGE_SIZE
)
torch.cuda.synchronize()
@@ -144,13 +146,13 @@ def test_q_rope_quant_matches_reference(pos_dtype):
_skip_if_unavailable()
dev = "cuda"
B, n_heads = 37, 64
cos, sin, freqs_cis, positions = _make_inputs(B, pos_dtype=pos_dtype)
cos, sin, cos_sin_cache, positions = _make_inputs(B, pos_dtype=pos_dtype)
q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev)
weight_scale = 0.137
q_fp8, weights_out = fused_q_indexer_rope_first_quant(
q, weight, weight_scale, freqs_cis, positions
q, weight, weight_scale, cos_sin_cache, positions
)
torch.cuda.synchronize()
@@ -181,7 +183,7 @@ def test_q_strided_weight_matches_contiguous():
_skip_if_unavailable()
dev = "cuda"
B, n_heads = 29, 64
cos, sin, freqs_cis, positions = _make_inputs(B)
cos, sin, cos_sin_cache, positions = _make_inputs(B)
q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
# weights_raw = kw[:, head_dim:] is a non-contiguous slice.
kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev)
@@ -190,16 +192,77 @@ def test_q_strided_weight_matches_contiguous():
assert not w_strided.is_contiguous()
a_fp8, a_w = fused_q_indexer_rope_first_quant(
q, w_strided, 0.137, freqs_cis, positions
q, w_strided, 0.137, cos_sin_cache, positions
)
b_fp8, b_w = fused_q_indexer_rope_first_quant(
q, w_contig, 0.137, freqs_cis, positions
q, w_contig, 0.137, cos_sin_cache, positions
)
torch.cuda.synchronize()
assert torch.equal(a_fp8, b_fp8)
assert torch.equal(a_w, b_w)
def test_indexer_uses_replaced_rope_cache_for_fused_kernels():
_skip_if_unavailable()
from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer
dev = "cuda"
B, n_heads = 7, 64
old_len = 16
new_len = 128
g = torch.Generator(device=dev).manual_seed(123)
old_cache = torch.randn(old_len, ROPE_DIM, dtype=torch.float32, device=dev)
cos = torch.randn(new_len, HALF, dtype=torch.float32, device=dev, generator=g)
sin = torch.randn(new_len, HALF, dtype=torch.float32, device=dev, generator=g)
grown_cache = torch.cat((cos, sin), dim=-1)
positions = torch.arange(old_len, old_len + B, device=dev, dtype=torch.int32)
class DummyRotary:
pass
rotary_emb = DummyRotary()
rotary_emb.cos_sin_cache = old_cache
indexer = Indexer.__new__(Indexer)
indexer.rotary_emb = rotary_emb
assert indexer._indexer_cos_sin_cache.data_ptr() == old_cache.data_ptr()
rotary_emb.cos_sin_cache = grown_cache
assert indexer._indexer_cos_sin_cache.data_ptr() == grown_cache.data_ptr()
key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev, generator=g)
k_weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g)
k_bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g)
k_out = fused_k_indexer_norm_rope(
key, k_weight, k_bias, EPS, indexer._indexer_cos_sin_cache, positions
)
normed = torch.nn.functional.layer_norm(
key.float(), (HEAD_DIM,), weight=k_weight, bias=k_bias, eps=EPS
)
k_ref = _rope_first(normed, cos[positions.long()], sin[positions.long()])
torch.testing.assert_close(k_out.float(), k_ref, atol=0.06, rtol=0.0)
q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev, generator=g)
q_weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev, generator=g)
weight_scale = 0.137
q_fp8, weights_out = fused_q_indexer_rope_first_quant(
q, q_weight, weight_scale, indexer._indexer_cos_sin_cache, positions
)
torch.cuda.synchronize()
q_ref = _rope_first(
q.float(), cos[positions.long()][:, None, :], sin[positions.long()][:, None, :]
)
scale = torch.clamp(q_ref.abs().amax(dim=-1, keepdim=True), min=1e-4) / FP8_MAX
torch.testing.assert_close(
weights_out.squeeze(-1),
q_weight.float() * weight_scale * scale.squeeze(-1),
atol=1e-3,
rtol=1e-3,
)
assert ((q_fp8.float() * scale - q_ref).abs() <= 0.0625 * q_ref.abs() + scale).all()
if __name__ == "__main__":
import sys