From 00d2d933a4e2036fa4c8680d04c88ca0ccf557fa Mon Sep 17 00:00:00 2001 From: Nam Le Date: Thu, 20 Aug 2026 04:42:49 +0000 Subject: [PATCH] [AMD] [GLM5] Fuse the DSA indexer Q/K rope + quant + K-cache write into one aiter kernel The DSA indexer Q/K fusion has been CUDA-only. On ROCm the indexer still runs the unfused path: rope on q, LayerNorm + rope on k, an fp8 quant of each, the head-gate scale, and the index-K cache write, as separate launches per layer per step. All of it is launch-bound. aiter's indexer_qk_rope_quant_and_cache does the whole thing in one launch, so extend the existing fusion switch to ROCm rather than build a second one: the fused wk_weights_proj GEMM, the no-Hadamard invariant, the weight loader and the graph split-op plumbing are all platform-agnostic already. - Probe aiter for the kernel at import and fall back with a warning; aiter is pinned per image and an older one would only fail at the first forward. - Build k_norm in fp32 whenever fusion is on. The kernel requires fp32 norm params, as the CUDA kernels already do; bf16 stays for the unfused ROCm path, where matching x.dtype is what selects aiter's CK layernorm. - Read cos/sin off aiter's rope module, which keeps them apart as [max_position, 1, 1, rope_dim/2] instead of one cos_sin_cache. - Give _fused_k_prepare_and_store a ROCm branch. Its CUDA fallback is a JIT kernel, and the k-only decode fast path reaches it whenever fusion is on. - Gate the fused-store branch on _is_cuda so non-CUDA stops paying for a JIT compile attempt that can only fail. Two adjacent fixes: - lora_manager imported _use_dsa_indexer_fusion, which #30111 deleted, so the indexer-LoRA guard raised ImportError instead of its intended error. Restored as dsa_indexer_fusion_supported(). - The ROCm branch of _store_index_k_cache read forward_batch.out_cache_loc, ignoring the sliced out_cache_loc its caller passes under the graph split-op contract. The two K-cache writers stay live at once -- the k-only decode path writes unfused -- and they are not byte-identical: the unfused path rounds to bf16 before quantizing while the fused kernel goes fp32 -> fp8 directly, so a value near an fp8 midpoint can land on either neighbour. Measured on 41x128: scales identical, 15/5248 elements differ, each by one fp8 code. Co-Authored-By: Claude Opus 5 (1M context) --- .../srt/layers/attention/dsa/dsa_indexer.py | 222 +++++++++++-- python/sglang/srt/lora/lora_manager.py | 4 +- .../test_dsa_indexer_qk_fuse_rocm.py | 294 ++++++++++++++++++ 3 files changed, 489 insertions(+), 31 deletions(-) create mode 100644 test/registered/kernels/ops/attention/test_dsa_indexer_qk_fuse_rocm.py diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index eaa5cc6044bd..60064c228ea6 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -108,6 +108,24 @@ if _use_aiter: from aiter.ops.cache import indexer_k_quant_and_cache + try: + from aiter.ops.cache import indexer_qk_rope_quant_and_cache + except ImportError: + indexer_qk_rope_quant_and_cache = None +else: + indexer_qk_rope_quant_and_cache = None + +# Single-launch aiter kernel covering Q rope + fp8 quant, K LayerNorm + rope + +# fp8 quant + index-K cache write, and the head-gate scale. aiter is pinned per +# image, so probe rather than assume: an older one imports fine and would only +# fail at the first layer's forward. +_use_aiter_indexer_qk_fuse = indexer_qk_rope_quant_and_cache is not None +if _use_aiter and not _use_aiter_indexer_qk_fuse: + logger.warning( + "ROCm DSA indexer: aiter has no indexer_qk_rope_quant_and_cache; " + "falling back to the unfused Q/K path." + ) + from sglang.srt.distributed import ( get_attn_tp_group, ) @@ -181,6 +199,17 @@ def _broadcast_indexer_topk_from_rank0( return topk_indices +def dsa_indexer_fusion_supported() -> bool: + """Whether this platform can run the fused indexer Q/K path. + + Indexer.__init__ adds the per-layer conditions; lora_manager reads this to + reject an indexer-targeted adapter, whose modules fusion folds away. + """ + return ( + _is_cuda or _use_aiter_indexer_qk_fuse + ) and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get() + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: # from sgl_kernel import hadamard_transform if _is_hip: @@ -235,10 +264,16 @@ def __init__( self.index_topk = index_topk self.q_lora_rank = q_lora_rank self.layer_id = layer_id + self.is_neox_style = is_neox_style + self._k_norm_is_rms = ( + config is not None + and getattr(config, "index_k_norm_type", "layer") == "rms" + ) + # The aiter kernel takes a LayerNorm weight and bias; RMSNorm has no bias. self.use_dsa_indexer_fusion = ( - _is_cuda - and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get() + dsa_indexer_fusion_supported() and not is_neox_style + and not (self._k_norm_is_rms and not _is_cuda) ) self.alt_stream = alt_stream self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() @@ -285,14 +320,19 @@ def __init__( params_dtype=torch.bfloat16, prefix=add_prefix("weights_proj", prefix), ) - if ( - config is not None - and getattr(config, "index_k_norm_type", "layer") == "rms" - ): + if self._k_norm_is_rms: self.k_norm = RMSNorm(self.head_dim) else: + # The fused kernels take fp32 norm params. Only the unfused ROCm path + # keeps bf16, where matching x.dtype is what selects aiter's CK + # layernorm over the native upcast. self.k_norm = LayerNorm( - self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 + self.head_dim, + dtype=( + torch.bfloat16 + if _use_aiter and not self.use_dsa_indexer_fusion + else torch.float32 + ), ) self.rotary_emb = get_rope_wrapper( rope_head_dim, @@ -334,6 +374,19 @@ def _with_real_sm_count(self): def _indexer_cos_sin_cache(self) -> torch.Tensor: return self.rotary_emb.cos_sin_cache + def _aiter_indexer_cos_sin(self) -> Tuple[torch.Tensor, torch.Tensor]: + """cos and sin as the [max_position, rope_head_dim // 2] pair aiter wants. + + aiter's rope module -- the one get_rope_wrapper builds under + SGLANG_USE_AITER -- keeps the two apart as [max_position, 1, 1, dim/2] + rather than in one cos_sin_cache, so this only drops the middle dims. + Both stay views; the kernel indexes them through cos_stride0 and needs + contiguity on the last dim alone. Read live rather than cached at + __init__, because the rope module's buffers can be replaced. + """ + cos, sin = self.rotary_emb.cos_cache, self.rotary_emb.sin_cache + return cos.view(cos.shape[0], -1), sin.view(sin.shape[0], -1) + def _weights_proj_bf16_in_fp32_out( self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] ) -> torch.Tensor: @@ -570,17 +623,39 @@ def _get_k_bf16( ): # Non-fusion path only; self.wk does not exist when fusion is on. key, _ = self.wk(x) - key = self.k_norm(key) + return rotate_activation(self._k_norm_rope(key, positions)) + + def _k_norm_rope( + self, key_raw: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + key = self.k_norm(key_raw) k_rope, _ = torch.split( key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - _, k_rope = self.rotary_emb(positions, k_rope, k_rope) self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) - key = rotate_activation(key) - return key + @staticmethod + def _acquire_index_k_pool(layer_id: int): + """The KV pool to write this layer's index-K into, or None if unowned.""" + pool = get_token_to_kv_pool() + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + if hasattr(pool, "_is_layer_owned") and not pool._is_layer_owned(layer_id): + return None + return pool + + @staticmethod + def _aiter_index_k_cache_view(pool, layer_id: int) -> torch.Tensor: + """[num_blocks, page_size, 132] fp8: 128 quantized elems + a 4-byte scale. + + The same view works for the legacy page_size=1 layout, where the middle + dim is 1. + """ + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + return buf.view(-1, pool.page_size, 132).view(fp8_dtype) + def _fused_k_prepare_and_store( self, key_raw: torch.Tensor, @@ -592,14 +667,24 @@ def _fused_k_prepare_and_store( ) -> None: if out_cache_loc is None: out_cache_loc = forward_batch.out_cache_loc - pool = get_token_to_kv_pool() + pool = self._acquire_index_k_pool(layer_id) + if pool is None: + return page_size = pool.page_size - if hasattr(pool, "invalidate_index_buffer_for_layer"): - pool.invalidate_index_buffer_for_layer(layer_id) - if hasattr(pool, "_is_layer_owned") and not pool._is_layer_owned(layer_id): + if _use_aiter: + # No Hadamard, matching _maybe_rotate under fusion, so this writes the + # same K representation the fused Q/K kernel and decode read back. + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=self._k_norm_rope(key_raw, positions), + act_quant=act_quant, + out_cache_loc=out_cache_loc, + ) return if ( - not _is_fp8_fnuz + _is_cuda + and not _is_fp8_fnuz and out_cache_loc is not None and can_use_dsa_fused_store(torch.bfloat16, out_cache_loc.dtype, page_size) ): @@ -647,6 +732,18 @@ def _fused_q_prepare_and_store( ) -> Tuple[torch.Tensor, torch.Tensor]: # num_tokens (graph split-op contract) slices q/k/positions/out_cache_loc # to the unpadded count; the returned q_fp8/weights are sliced to match. + if _use_aiter_indexer_qk_fuse: + # One kernel covers Q and K, so there is nothing for the dual-stream + # split below to overlap. + return self._aiter_fused_qk_prepare_and_store( + x=x, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + act_quant=act_quant, + num_tokens=num_tokens, + ) q_scale_gate = self.softmax_scale * self.n_heads**-0.5 out_cache_loc = forward_batch.out_cache_loc if num_tokens is not None: @@ -717,6 +814,82 @@ def _fused_q_prepare_and_store( current_stream.wait_stream(self.alt_stream) return q_fp8, weights + def _aiter_fused_qk_prepare_and_store( + self, + *, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + num_tokens: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + out_cache_loc = forward_batch.out_cache_loc + if num_tokens is not None: + positions = positions[:num_tokens] + out_cache_loc = out_cache_loc[:num_tokens] + + key, weights_raw = self._fused_k_weights(x) + q = self.wq_b(q_lora)[0].view(-1, self.n_heads, self.head_dim) + if num_tokens is not None: + key = key[:num_tokens] + weights_raw = weights_raw[:num_tokens] + q = q[:num_tokens] + + pool = self._acquire_index_k_pool(layer_id) + if pool is None: + return self._q_rope_quant_and_gate(q, weights_raw, positions, act_quant) + + if not out_cache_loc.is_contiguous(): + out_cache_loc = out_cache_loc.contiguous() + cos, sin = self._aiter_indexer_cos_sin() + q_fp8 = torch.empty(q.shape, dtype=fp8_dtype, device=q.device) + weights = torch.empty( + (q.shape[0], self.n_heads), dtype=torch.float32, device=q.device + ) + # key and weights_raw stay strided views of the wk_weights_proj output; + # the kernel indexes both through their strides. + indexer_qk_rope_quant_and_cache( + q, + q_fp8, + weights_raw, + weights, + key, + self._aiter_index_k_cache_view(pool, layer_id), + out_cache_loc, + self.k_norm.weight, + self.k_norm.bias, + positions, + cos, + sin, + self.k_norm.variance_epsilon, + self.block_size, + self.scale_fmt, + self.softmax_scale * self.n_heads**-0.5, + preshuffle=_use_aiter_preshuffle, + is_neox=self.is_neox_style, + ) + # [num_tokens, n_heads, 1] like the CUDA fused Q kernel; the top-k paths + # assert rank 3 and squeeze. + return q_fp8, weights.unsqueeze(-1) + + def _q_rope_quant_and_gate( + self, + q: torch.Tensor, + weights_raw: torch.Tensor, + positions: torch.Tensor, + act_quant, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Q half of the fused kernel, for a layer whose index-K another rank owns.""" + q_rope, _ = torch.split( + q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + q_rope, _ = self.rotary_emb(positions, q_rope, q_rope) + self._update_rope_guarded(q[..., : self.rope_head_dim], q_rope) + q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) + return q_fp8, self._scale_head_gates(weights_raw, q_scale) + @staticmethod def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None: # On AMD with in-place RoPE kernels, self-aliasing can occur; @@ -1474,10 +1647,8 @@ def _store_index_k_cache( if out_cache_loc is None: out_cache_loc = forward_batch.out_cache_loc - pool = get_token_to_kv_pool() - if hasattr(pool, "invalidate_index_buffer_for_layer"): - pool.invalidate_index_buffer_for_layer(layer_id) - if hasattr(pool, "_is_layer_owned") and not pool._is_layer_owned(layer_id): + pool = self._acquire_index_k_pool(layer_id) + if pool is None: return if ( @@ -1500,20 +1671,13 @@ def _store_index_k_cache( return # Fast path: AITER fused quant + cache store - # When _use_aiter_preshuffle is True we use the new MFMA 16x16 preshuffle - # layout (page_size>=16). Otherwise we fall back to the legacy row-major - # layout with page_size=1; the same kv_cache.view works for both cases - # because page_size is 1 there. if _use_aiter: - page_size = pool.page_size - buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) - kv_cache = buf.view(-1, page_size, 132).view(fp8_dtype) - out_loc = forward_batch.out_cache_loc + out_loc = out_cache_loc if not out_loc.is_contiguous(): out_loc = out_loc.contiguous() indexer_k_quant_and_cache( key, - kv_cache, + self._aiter_index_k_cache_view(pool, layer_id), out_loc, self.block_size, self.scale_fmt, diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index aae2afffc7d4..6f443a566f3f 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -732,10 +732,10 @@ def init_lora_shapes( indexer_targets = self.target_modules & DSA_INDEXER_LORA_NAMES if indexer_targets: from sglang.srt.layers.attention.dsa.dsa_indexer import ( - _use_dsa_indexer_fusion, + dsa_indexer_fusion_supported, ) - if _use_dsa_indexer_fusion: + if dsa_indexer_fusion_supported(): raise ValueError( f"LoRA targets the DSA indexer ({sorted(indexer_targets)}), which is " "incompatible with DSA indexer Q/K fusion. Set " diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer_qk_fuse_rocm.py b/test/registered/kernels/ops/attention/test_dsa_indexer_qk_fuse_rocm.py new file mode 100644 index 000000000000..9f372d730fbc --- /dev/null +++ b/test/registered/kernels/ops/attention/test_dsa_indexer_qk_fuse_rocm.py @@ -0,0 +1,294 @@ +"""Correctness tests for the ROCm/aiter DSA indexer Q/K fusion. + +``indexer_qk_rope_quant_and_cache`` replaces five launches per layer -- Q rope, +K LayerNorm, K rope, fp8 quant of both, and the index-K cache write -- plus the +head-gate scale. Every case here pins that single kernel against the unfused +ROCm path it replaces, because the two must stay interchangeable: the k-only +decode fast path still writes the cache the unfused way, and decode reads back +whatever either wrote. +""" + +from __future__ import annotations + +import pytest +import torch + +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=45, suite="jit-kernel-unit-test-amd") + +HEAD_DIM = 128 +N_HEADS = 32 +ROPE_DIM = 64 +HALF = ROPE_DIM // 2 +PAGE_SIZE = 64 +BYTES_PER_TOKEN = HEAD_DIM + 4 # 128 fp8 + 4-byte fp32 scale +CACHE_STRIDE = BYTES_PER_TOKEN +EPS = 1e-5 +MAX_POS = 8192 +BLOCK_SIZE = 128 +SCALE_FMT = "ue8m0" +WEIGHTS_SCALE = HEAD_DIM**-0.5 * N_HEADS**-0.5 + + +def _skip_if_unavailable(): + if not is_hip(): + pytest.skip("aiter indexer Q/K fusion is ROCm-specific") + if not torch.cuda.is_available(): + pytest.skip("GPU required") + pytest.importorskip("aiter") + from aiter.ops import cache as aiter_cache + + if not hasattr(aiter_cache, "indexer_qk_rope_quant_and_cache"): + pytest.skip("aiter lacks indexer_qk_rope_quant_and_cache") + + +def _aiter_ops(): + from aiter.ops.cache import ( + indexer_k_quant_and_cache, + indexer_qk_rope_quant_and_cache, + ) + + return indexer_qk_rope_quant_and_cache, indexer_k_quant_and_cache + + +def _fp8_dtype(): + from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype + + return fp8_dtype + + +def _make_inputs(B, n_heads=N_HEADS, seed=0, strided=False): + g = torch.Generator(device="cuda").manual_seed(seed) + dev = "cuda" + cos = torch.randn(MAX_POS, HALF, dtype=torch.bfloat16, device=dev, generator=g) + sin = torch.randn(MAX_POS, HALF, dtype=torch.bfloat16, device=dev, generator=g) + positions = torch.randint(0, 4096, (B,), device=dev, dtype=torch.int64, generator=g) + q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev, generator=g) + # The real inputs are slices of one wk_weights_proj GEMM output, so they are + # strided; the fused path passes them through without a contiguous copy. + kw = torch.randn( + B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev, generator=g + ) + key, weights_raw = kw[:, :HEAD_DIM], kw[:, HEAD_DIM:] + if not strided: + key, weights_raw = key.contiguous(), weights_raw.contiguous() + # The kernel requires fp32 norm params (cache_kernels.cu), which is why the + # Indexer builds k_norm in fp32 whenever fusion is on. + norm_weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g) + norm_bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g) + return cos, sin, positions, q, key, weights_raw, norm_weight, norm_bias + + +def _make_cache(B, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed + 991) + loc = torch.randperm(B * 4, device="cuda", generator=g)[:B].to(torch.int64) + num_pages = int(loc.max().item()) // PAGE_SIZE + 2 + buf = torch.zeros( + num_pages, CACHE_STRIDE * PAGE_SIZE, dtype=torch.uint8, device="cuda" + ) + return buf, loc, num_pages + + +def _cache_view(buf, page_size=PAGE_SIZE): + return buf.view(-1, page_size, CACHE_STRIDE).view(_fp8_dtype()) + + +def _rope_interleaved(x, cos_p, sin_p): + """Interleaved (indexer_rope_interleave) rope on the leading ROPE_DIM dims.""" + x = x.clone() + xr = x[..., 0:ROPE_DIM:2].clone() + xi = x[..., 1:ROPE_DIM:2].clone() + x[..., 0:ROPE_DIM:2] = xr * cos_p - xi * sin_p + x[..., 1:ROPE_DIM:2] = xr * sin_p + xi * cos_p + return x + + +def _run_fused(inputs, buf, loc, preshuffle, page_size=PAGE_SIZE): + fused, _ = _aiter_ops() + cos, sin, positions, q, key, weights_raw, norm_weight, norm_bias = inputs + q_fp8 = torch.empty(q.shape, dtype=_fp8_dtype(), device=q.device) + weights = torch.empty( + (q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device + ) + fused( + q, + q_fp8, + weights_raw, + weights, + key, + _cache_view(buf, page_size), + loc, + norm_weight, + norm_bias, + positions, + cos, + sin, + EPS, + BLOCK_SIZE, + SCALE_FMT, + WEIGHTS_SCALE, + preshuffle=preshuffle, + is_neox=False, + ) + torch.cuda.synchronize() + return q_fp8, weights + + +def _unfused_k_cache(inputs, loc, npages, preshuffle, page_size): + """What the k-only decode path writes: bf16 LayerNorm+rope, then a quant kernel.""" + _, k_quant_and_cache = _aiter_ops() + cos, sin, positions, _, key, _, norm_weight, norm_bias = inputs + normed = torch.nn.functional.layer_norm( + key.float(), (HEAD_DIM,), weight=norm_weight, bias=norm_bias, eps=EPS + ) + key_bf16 = _rope_interleaved( + normed, cos[positions].float(), sin[positions].float() + ).to(torch.bfloat16) + buf = torch.zeros( + npages, CACHE_STRIDE * page_size, dtype=torch.uint8, device="cuda" + ) + k_quant_and_cache( + key_bf16, + buf.view(-1, page_size, CACHE_STRIDE).view(_fp8_dtype()), + loc, + BLOCK_SIZE, + SCALE_FMT, + preshuffle=preshuffle, + ) + torch.cuda.synchronize() + return buf + + +@pytest.mark.parametrize("preshuffle", [False, True]) +def test_k_cache_agrees_with_the_unfused_writer(preshuffle): + """Both writers stay live -- the k-only decode path uses the unfused one -- + and decode reads back whichever wrote the token, so they must agree. + + Not byte-identical by construction: the unfused path rounds to bf16 before + quantizing, the fused kernel goes fp32 -> fp8 in one step, so a value near an + fp8 midpoint can land on either neighbour. + """ + _skip_if_unavailable() + B = 41 + inputs = _make_inputs(B) + buf_fused, loc, npages = _make_cache(B) + _run_fused(inputs, buf_fused, loc, preshuffle) + buf_unfused = _unfused_k_cache(inputs, loc, npages, preshuffle, PAGE_SIZE) + + differing = int((buf_fused != buf_unfused).sum()) + assert differing < 0.01 * buf_fused.numel(), ( + f"{differing}/{buf_fused.numel()} bytes differ; double rounding alone " + "moves far fewer than 1%" + ) + + +def test_k_cache_is_within_one_fp8_step_of_the_unfused_writer(): + """The flat page_size=1 layout is the one this test can decode, so the exact + numeric bound lives here and the paged layouts get the byte-fraction check.""" + _skip_if_unavailable() + B = 41 + page_size = 1 + inputs = _make_inputs(B, seed=11) + loc = torch.arange(B, device="cuda", dtype=torch.int64) + npages = B + 2 + buf_fused = torch.zeros( + npages, CACHE_STRIDE * page_size, dtype=torch.uint8, device="cuda" + ) + _run_fused(inputs, buf_fused, loc, preshuffle=False, page_size=page_size) + buf_unfused = _unfused_k_cache(inputs, loc, npages, False, page_size) + + def decode(buf): + v = buf.view(-1, CACHE_STRIDE)[:B] + payload = v[:, :HEAD_DIM].contiguous().view(_fp8_dtype()).float() + scale = v[:, HEAD_DIM:].contiguous().view(torch.float32) + return payload * scale, scale + + fused, scale_fused = decode(buf_fused) + unfused, scale_unfused = decode(buf_unfused) + + # A differing ue8m0 exponent would rescale a whole token, which is a bigger + # claim than double rounding can make. + assert torch.equal(scale_fused, scale_unfused) + # Adjacent fp8-e4m3 codes are at most 1/8 apart in relative terms (3 mantissa + # bits, worst case at the bottom of a binade); one scale step covers zero. + assert ( + (fused - unfused).abs() <= 0.125 * unfused.abs() + scale_unfused + ).all(), f"max deviation {(fused - unfused).abs().max().item()}" + + +def test_q_and_head_gate_match_unfused(): + """q_fp8 and the folded head gate must match rope + act_quant + _scale_head_gates.""" + _skip_if_unavailable() + from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant + + B = 37 + inputs = _make_inputs(B, seed=3) + cos, sin, positions, q, _, weights_raw, _, _ = inputs + buf, loc, _ = _make_cache(B, seed=3) + q_fp8, weights = _run_fused(inputs, buf, loc, preshuffle=True) + + cp = cos[positions].float()[:, None, :] + sp = sin[positions].float()[:, None, :] + q_roped = _rope_interleaved(q.float(), cp, sp).to(torch.bfloat16) + q_fp8_ref, q_scale_ref = act_quant(q_roped, BLOCK_SIZE, SCALE_FMT) + torch.cuda.synchronize() + + # ue8m0 rounds the scale to a power of two, so both paths must land on the + # identical scale; only the fp8 payload may differ by a rounding step. + weights_ref = weights_raw.float() * WEIGHTS_SCALE * q_scale_ref.squeeze(-1) + torch.testing.assert_close(weights, weights_ref, atol=0, rtol=1e-6) + + deq = q_fp8.float() * q_scale_ref + deq_ref = q_fp8_ref.float() * q_scale_ref + # fp8-e4m3 has 3 mantissa bits: one rounding step is 1/16 relative. + err = (deq - deq_ref).abs() + assert ( + err <= 0.0625 * deq_ref.abs() + q_scale_ref + ).all(), f"max fp8 mismatch {err.max().item()}" + + +def test_strided_inputs_match_contiguous(): + """The no-copy path: key/weights_raw as wk_weights_proj slices, not copies.""" + _skip_if_unavailable() + B = 29 + strided = _make_inputs(B, seed=7, strided=True) + contig = _make_inputs(B, seed=7, strided=False) + assert not strided[4].is_contiguous() and not strided[5].is_contiguous() + + buf_a, loc, _ = _make_cache(B, seed=7) + buf_b = torch.zeros_like(buf_a) + q_fp8_a, w_a = _run_fused(strided, buf_a, loc, preshuffle=True) + q_fp8_b, w_b = _run_fused(contig, buf_b, loc, preshuffle=True) + + assert torch.equal(buf_a, buf_b) + assert torch.equal(q_fp8_a, q_fp8_b) + assert torch.equal(w_a, w_b) + + +def test_cos_sin_view_tracks_a_replaced_rope_cache(): + """aiter keeps cos/sin as [max_position, 1, 1, dim/2]; the kernel needs 2-D. + + Read live, not cached at __init__, so a rope module whose buffers were + replaced (a grown cache) is picked up. + """ + from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer + + class DummyRotary: + pass + + indexer = Indexer.__new__(Indexer) + indexer.rotary_emb = DummyRotary() + old_cos = torch.randn(16, 1, 1, HALF, dtype=torch.bfloat16) + indexer.rotary_emb.cos_cache = old_cos + indexer.rotary_emb.sin_cache = torch.randn(16, 1, 1, HALF, dtype=torch.bfloat16) + cos, sin = indexer._aiter_indexer_cos_sin() + assert cos.shape == (16, HALF) and sin.shape == (16, HALF) + assert cos.data_ptr() == old_cos.data_ptr() + + grown = torch.randn(128, 1, 1, HALF, dtype=torch.bfloat16) + indexer.rotary_emb.cos_cache = grown + indexer.rotary_emb.sin_cache = torch.randn(128, 1, 1, HALF, dtype=torch.bfloat16) + cos, _ = indexer._aiter_indexer_cos_sin() + assert cos.shape == (128, HALF) and cos.data_ptr() == grown.data_ptr()