diff --git a/aphrodite/models/deepseek_v4/attention.py b/aphrodite/models/deepseek_v4/attention.py index 7c987c5da4..594956e684 100644 --- a/aphrodite/models/deepseek_v4/attention.py +++ b/aphrodite/models/deepseek_v4/attention.py @@ -29,6 +29,7 @@ from aphrodite.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE if TYPE_CHECKING: + from aphrodite.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from aphrodite.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) @@ -181,6 +182,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None = None, q_workspace: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ) -> None: super().__init__() config = aphrodite_config.model_config.hf_config @@ -270,6 +272,7 @@ def __init__( self.indexer_rotary_emb = self.rotary_emb self.topk_indices_buffer = topk_indices_buffer self.q_workspace = q_workspace + self.eager_scratch_pool = eager_scratch_pool self.indexer = None if self.compress_ratio == 4: @@ -289,6 +292,7 @@ def __init__( compress_ratio=self.compress_ratio, prefix=f"{prefix}.indexer", aux_stream=indexer_aux_stream, + eager_scratch_pool=eager_scratch_pool, ) # Will be None on ROCm for now. @@ -337,6 +341,7 @@ def __init__( rotate=True, prefix=f"{prefix}.compressor", k_cache_prefix=self.prefix, + eager_scratch_pool=eager_scratch_pool, ) def forward( @@ -565,7 +570,9 @@ def _fused_qnorm_rope_kv_insert( # the padded q tensor. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - if self.q_workspace is None: + if self.eager_scratch_pool is not None: + q_out = self.eager_scratch_pool.q_out(q.shape[0]) + elif self.q_workspace is None: q_out = torch.empty( (q.shape[0], self.padded_heads, self.head_dim), dtype=q.dtype, @@ -622,6 +629,11 @@ def _fused_qnorm_rope_kv_insert( ) return q_fp8 + def _global_topk_output_buffers(self, topk_indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] | None: + if self.compress_ratio != 4 or self.eager_scratch_pool is None: + return None + return self.eager_scratch_pool.global_topk_outputs(topk_indices) + def get_attn_backend(self) -> type[AttentionBackend]: return self.backend_cls @@ -699,6 +711,7 @@ def __init__( compress_ratio: int = 1, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.aphrodite_config = aphrodite_config @@ -711,6 +724,7 @@ def __init__( self.rope_dim = config.qk_rope_head_dim # 64 self.q_lora_rank = q_lora_rank # 1536 self.compress_ratio = compress_ratio + self.eager_scratch_pool = eager_scratch_pool self.use_fp4_kv = self.aphrodite_config.attention_config.use_fp4_indexer_cache logger.info_once( "Using %s indexer cache for Lightning Indexer.", @@ -768,6 +782,7 @@ def __init__( prefix=f"{prefix}.compressor", k_cache_prefix=self.k_cache.prefix, use_fp4_cache=self.use_fp4_kv, + eager_scratch_pool=eager_scratch_pool, ) self.indexer_op = SparseAttnIndexer( @@ -822,6 +837,9 @@ def wq_b_and_q_quant(): # ReplicatedLinear returns (output, bias); bias is None. q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) + outputs = None + if self.eager_scratch_pool is not None and self.use_fp4_kv: + outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0]) return fused_indexer_q_rope_quant( positions, q, @@ -830,6 +848,7 @@ def wq_b_and_q_quant(): self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + output_buffers=outputs, ) # compressor returns None and writes K to the indexer KV cache; the diff --git a/aphrodite/models/deepseek_v4/common/ops/cache_utils.py b/aphrodite/models/deepseek_v4/common/ops/cache_utils.py index 30eb2a47ee..9b95c56f4a 100644 --- a/aphrodite/models/deepseek_v4/common/ops/cache_utils.py +++ b/aphrodite/models/deepseek_v4/common/ops/cache_utils.py @@ -431,6 +431,7 @@ def compute_global_topk_indices_and_lens( block_table: torch.Tensor, block_size: int, is_valid_token: torch.Tensor, + output_buffers: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Map local topk indices to global KV cache slots and count valid entries. @@ -440,8 +441,13 @@ def compute_global_topk_indices_and_lens( 3. Masking padding tokens to length 0 """ num_tokens = topk_indices.shape[0] - global_topk_indices = torch.empty_like(topk_indices) - topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + if output_buffers is None: + global_topk_indices = torch.empty_like(topk_indices) + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + else: + global_topk_indices, topk_lens = output_buffers + assert global_topk_indices.shape == topk_indices.shape + assert topk_lens.shape == (num_tokens,) _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( global_topk_indices, global_topk_indices.stride(0), diff --git a/aphrodite/models/deepseek_v4/common/ops/fused_indexer_q.py b/aphrodite/models/deepseek_v4/common/ops/fused_indexer_q.py index 4c8dae062a..87e93c4ed2 100644 --- a/aphrodite/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/aphrodite/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -276,6 +276,7 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + output_buffers: tuple[torch.Tensor, ...] | None = None, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, @@ -313,23 +314,36 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + if output_buffers is None: + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + else: + expected_num_buffers = 3 if use_fp4 else 2 + assert len(output_buffers) == expected_num_buffers + index_weights_out = output_buffers[-1] + assert index_weights_out.shape == index_weights.shape if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block size {MXFP4_BLOCK_SIZE}" ) num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE - index_q_packed = torch.empty( - (num_tokens, num_index_q_heads, index_q_head_dim // 2), - dtype=torch.uint8, - device=index_q.device, - ) - index_q_scale = torch.empty( - (num_tokens, num_index_q_heads, num_scale_blocks), - dtype=torch.uint8, - device=index_q.device, - ) + packed_shape = (num_tokens, num_index_q_heads, index_q_head_dim // 2) + scale_shape = (num_tokens, num_index_q_heads, num_scale_blocks) + if output_buffers is None: + index_q_packed = torch.empty( + packed_shape, + dtype=torch.uint8, + device=index_q.device, + ) + index_q_scale = torch.empty( + scale_shape, + dtype=torch.uint8, + device=index_q.device, + ) + else: + index_q_packed, index_q_scale, _ = output_buffers + assert index_q_packed.shape == packed_shape + assert index_q_scale.shape == scale_shape if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from aphrodite.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( @@ -398,7 +412,11 @@ def fused_indexer_q_rope_quant( fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 - index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) + if output_buffers is None: + index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) + else: + index_q_fp8, _ = output_buffers + assert index_q_fp8.shape == index_q.shape if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from aphrodite.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( diff --git a/aphrodite/models/deepseek_v4/compressor.py b/aphrodite/models/deepseek_v4/compressor.py index 5b126e81d2..6fe4f97637 100644 --- a/aphrodite/models/deepseek_v4/compressor.py +++ b/aphrodite/models/deepseek_v4/compressor.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import torch from torch import nn @@ -39,6 +39,9 @@ SlidingWindowMLASpec, ) +if TYPE_CHECKING: + from aphrodite.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool + def _prefer_two_stage_compressor() -> bool: # Platforms that favor the triton variant of two-stage compressor split. @@ -219,6 +222,7 @@ def __init__( prefix: str = "", k_cache_prefix="", use_fp4_cache: bool = False, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.compress_ratio = compress_ratio @@ -228,6 +232,7 @@ def __init__( self.prefix = prefix self.k_cache_prefix = k_cache_prefix self.use_fp4_cache = use_fp4_cache + self.eager_scratch_pool = eager_scratch_pool config = aphrodite_config.model_config.hf_config self.rope_head_dim = config.qk_rope_head_dim @@ -399,6 +404,8 @@ def forward( store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) + if not self.overlap and self.eager_scratch_pool is not None: + extra_kwargs["compress_scratch"] = self.eager_scratch_pool.compressor_scratch(num_actual) elif self._use_two_stage_fused_compressor: # head=512 cr>=128 (no overlap): two-pass split compressor on the # prefill suffix, single-pass on the decode prefix. diff --git a/aphrodite/models/deepseek_v4/eager_scratch.py b/aphrodite/models/deepseek_v4/eager_scratch.py new file mode 100644 index 0000000000..48de4dc115 --- /dev/null +++ b/aphrodite/models/deepseek_v4/eager_scratch.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from math import prod + +import torch + +from aphrodite.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE +from aphrodite.utils.math_utils import round_up + + +class DeepseekV4EagerScratchPool: + """Model-wide outputs and scratch used inside the attention eager break.""" + + _ALIGNMENT = 256 + + def __init__( + self, + max_num_tokens: int, + q_workspace: torch.Tensor, + q_head_dim: int, + index_q_heads: int, + index_q_head_dim: int, + index_topk: int, + device: torch.device | str, + ) -> None: + self.max_num_tokens = max_num_tokens + self.index_topk = index_topk + assert q_workspace.shape[0] >= max_num_tokens + assert q_workspace.shape[-1] == q_head_dim + self._q = q_workspace + + fp4_specs = ( + ((max_num_tokens, index_q_heads, index_q_head_dim // 2), torch.uint8), + ( + ( + max_num_tokens, + index_q_heads, + index_q_head_dim // MXFP4_BLOCK_SIZE, + ), + torch.uint8, + ), + ((max_num_tokens, index_q_heads), torch.float32), + ) + global_specs = ( + ((max_num_tokens, index_topk), torch.int32), + ((max_num_tokens,), torch.int32), + ) + compressor_specs = (((max_num_tokens, q_head_dim), torch.float32),) + # FP4 indexer is C4 only, global mapping after FP4 indexer + # compressor scratch is C128 only + # so here we use max instead of sum + aux_bytes = max(self._packed_size(specs) for specs in (fp4_specs, global_specs, compressor_specs)) + storage = torch.empty(aux_bytes, dtype=torch.uint8, device=device) + + self._q_outputs: dict[int, torch.Tensor] = {} + fp4_values, fp4_scales, fp4_weights = self._views(storage, fp4_specs) + self._fp4_template = (fp4_values, fp4_scales, fp4_weights) + self._fp4_outputs: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + global_indices, global_lens = self._views(storage, global_specs) + self._global_template = (global_indices, global_lens) + self._global_outputs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + self._compressor_template = self._views(storage, compressor_specs)[0] + self._compressor_outputs: dict[int, torch.Tensor] = {} + self._storage = storage + + @classmethod + def _packed_size(cls, specs: tuple[tuple[tuple[int, ...], torch.dtype], ...]) -> int: + offset = 0 + for shape, dtype in specs: + offset = round_up(offset, cls._ALIGNMENT) + prod(shape) * dtype.itemsize + return round_up(offset, cls._ALIGNMENT) + + @classmethod + def _views( + cls, + storage: torch.Tensor, + specs: tuple[tuple[tuple[int, ...], torch.dtype], ...], + ) -> list[torch.Tensor]: + offset = 0 + views = [] + for shape, dtype in specs: + offset = round_up(offset, cls._ALIGNMENT) + num_bytes = prod(shape) * dtype.itemsize + views.append(storage[offset : offset + num_bytes].view(dtype).view(shape)) + offset += num_bytes + return views + + def q_out(self, num_tokens: int) -> torch.Tensor: + output = self._q_outputs.get(num_tokens) + if output is None: + output = self._q[:num_tokens] + self._q_outputs[num_tokens] = output + return output + + def compressor_scratch(self, num_tokens: int) -> torch.Tensor: + output = self._compressor_outputs.get(num_tokens) + if output is None: + output = self._compressor_template[:num_tokens] + self._compressor_outputs[num_tokens] = output + return output + + def indexer_q_outputs( + self, + num_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + output = self._fp4_outputs.get(num_tokens) + if output is None: + values, scales, weights = self._fp4_template + output = ( + values[:num_tokens], + scales[:num_tokens], + weights[:num_tokens], + ) + self._fp4_outputs[num_tokens] = output + return output + + def global_topk_outputs(self, topk_indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens, topk = topk_indices.shape + assert topk == self.index_topk + output = self._global_outputs.get(num_tokens) + if output is None: + indices, lens = self._global_template + output = (indices[:num_tokens], lens[:num_tokens]) + self._global_outputs[num_tokens] = output + return output diff --git a/aphrodite/models/deepseek_v4/nvidia/dspark.py b/aphrodite/models/deepseek_v4/nvidia/dspark.py index ff4422c439..675497c2b0 100644 --- a/aphrodite/models/deepseek_v4/nvidia/dspark.py +++ b/aphrodite/models/deepseek_v4/nvidia/dspark.py @@ -86,12 +86,19 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> No ) self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.topk_indices_buffer = torch.empty( + aphrodite_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) + current_aphrodite_config = get_current_aphrodite_config() self.layers = nn.ModuleList( [ DeepseekV4DecoderLayer( current_aphrodite_config, prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + topk_indices_buffer=self.topk_indices_buffer, ) for i in range(self.num_dspark_layers) ] diff --git a/aphrodite/models/deepseek_v4/nvidia/flashinfer_sparse.py b/aphrodite/models/deepseek_v4/nvidia/flashinfer_sparse.py index f03fb46ec0..60795fe9be 100644 --- a/aphrodite/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/aphrodite/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -694,6 +694,7 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, + output_buffers=self._global_topk_output_buffers(self.topk_indices_buffer[:num_decode_tokens]), ) extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) else: @@ -773,6 +774,7 @@ def _forward_prefill( attn_metadata.block_table, block_size, swa_metadata.is_valid_token[prefill_token_slice], + output_buffers=self._global_topk_output_buffers(local_topk_indices), ) assert swa_metadata.prefill_swa_indices is not None diff --git a/aphrodite/models/deepseek_v4/nvidia/flashmla.py b/aphrodite/models/deepseek_v4/nvidia/flashmla.py index 57dabac2af..c0e0aa33c1 100644 --- a/aphrodite/models/deepseek_v4/nvidia/flashmla.py +++ b/aphrodite/models/deepseek_v4/nvidia/flashmla.py @@ -165,6 +165,7 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, + output_buffers=self._global_topk_output_buffers(self.topk_indices_buffer[:num_decode_tokens]), ) topk_indices = global_indices.view(num_decode_tokens, 1, -1) else: diff --git a/aphrodite/models/deepseek_v4/nvidia/model.py b/aphrodite/models/deepseek_v4/nvidia/model.py index 715e352a6d..f78bcffdd9 100644 --- a/aphrodite/models/deepseek_v4/nvidia/model.py +++ b/aphrodite/models/deepseek_v4/nvidia/model.py @@ -66,6 +66,7 @@ ) from aphrodite.model_executor.utils import set_weight_attrs from aphrodite.models.deepseek_v4.attention import DeepseekV4Attention +from aphrodite.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from aphrodite.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -774,6 +775,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None = None, q_workspace: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + eager_scratch_pool: DeepseekV4EagerScratchPool | None = None, ): super().__init__() @@ -787,6 +789,7 @@ def __init__( topk_indices_buffer=topk_indices_buffer, q_workspace=q_workspace, aux_stream_list=aux_stream_list, + eager_scratch_pool=eager_scratch_pool, ) self.ffn = DeepseekV4MoE(aphrodite_config, prefix=f"{prefix}.ffn") @@ -976,6 +979,18 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): config.head_dim, dtype=aphrodite_config.model_config.dtype, ) + self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None + if not aphrodite_config.parallel_config.use_ubatching: + # TODO: support DBO if needed; this requires a microbatch dimension. + self.eager_scratch_pool = DeepseekV4EagerScratchPool( + aphrodite_config.scheduler_config.max_num_batched_tokens, + self.q_workspace, + config.head_dim, + config.index_n_heads, + config.index_head_dim, + config.index_topk, + current_platform.device_type, + ) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -995,6 +1010,7 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): topk_indices_buffer=self.topk_indices_buffer, q_workspace=self.q_workspace, aux_stream_list=aux_stream_list, + eager_scratch_pool=self.eager_scratch_pool, ), prefix=f"{prefix}.layers", ) diff --git a/aphrodite/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/aphrodite/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index f04376bdf6..8642489890 100644 --- a/aphrodite/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/aphrodite/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -1942,6 +1942,7 @@ def compress_norm_rope_store_cutedsl( store_full_kv: bool = False, store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, + compress_scratch: torch.Tensor | None = None, ) -> None: if compress_ratio == 4: # For C4A, the single fused kernel is faster than the two-kernel version. @@ -1974,11 +1975,15 @@ def compress_norm_rope_store_cutedsl( ) else: # For C128, the two-kernel version is faster than the single fused kernel. - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, - ) + if compress_scratch is None: + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) + else: + assert compress_scratch.shape == (num_actual, head_dim) + compressed_kv = compress_scratch split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( state_cache, token_to_req_indices, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cd75fe3185..e7d1f86f6c 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -20,6 +20,7 @@ from aphrodite import _custom_ops as ops from aphrodite.models.deepseek_v4.common.ops import ( + compute_global_topk_indices_and_lens, dequantize_and_gather_k_cache, quantize_and_insert_k_cache, ) @@ -34,6 +35,21 @@ from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +def test_compute_global_topk_reuses_output_buffers(): + device = "cuda" + topk_indices = torch.tensor([[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device) + token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device) + block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device) + is_valid = torch.tensor([True, False], device=device) + args = (topk_indices, token_to_req, block_table, 2, is_valid) + expected = compute_global_topk_indices_and_lens(*args) + outputs = tuple(torch.empty_like(tensor) for tensor in expected) + actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs) + for result, output, reference in zip(actual, outputs, expected): + assert result.data_ptr() == output.data_ptr() + torch.testing.assert_close(result, reference) + + def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale). diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index dda0f8479e..0d1bf26f4b 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -144,6 +144,21 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use head_scale = N_HEAD**-0.5 q_quant_ref, weights_ref = _reference(positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4) + output_buffers: tuple[torch.Tensor, ...] | None = None + output_buffer_test_num_tokens = 7 + if num_tokens == output_buffer_test_num_tokens and cache_dtype == torch.float32: + if use_fp4: + q_ref, q_scale_ref = q_quant_ref + output_buffers = ( + torch.empty_like(q_ref), + torch.empty_like(q_scale_ref).view(torch.uint8).reshape(num_tokens, N_HEAD, -1), + torch.empty_like(weights_ref), + ) + else: + output_buffers = ( + torch.empty_like(q_quant_ref), + torch.empty_like(weights_ref), + ) # use_cutedsl=False: force the triton path even when cutedsl is installed # by patching the dispatcher's has_cutedsl() binding to return False. cutedsl_patch = ( @@ -163,8 +178,17 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use softmax_scale, head_scale, use_fp4, + output_buffers=output_buffers, ) + if output_buffers is not None: + if use_fp4: + assert q_quant_fused[0].data_ptr() == output_buffers[0].data_ptr() + assert q_quant_fused[1].data_ptr() == output_buffers[1].data_ptr() + else: + assert q_quant_fused.data_ptr() == output_buffers[0].data_ptr() + assert weights_fused.data_ptr() == output_buffers[-1].data_ptr() + if use_fp4: q_quant_ref, q_scale_ref = q_quant_ref q_quant_fused, q_scale_fused = q_quant_fused