minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 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
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