diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 05cb2cbe517a..9288cecf9ed5 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -186,11 +186,13 @@ def _handle_dflash(server_args: ServerArgs) -> None: cfg = resolving_view(server_args) from sglang.srt.arg_groups.overrides import resolved_view - if not (cfg.device.startswith("cuda") or cfg.device == "npu"): + if not (cfg.device.startswith("cuda") or cfg.device in ("npu", "cpu")): raise ValueError( - "DFLASH speculative decoding only supports CUDA and NPU devices." + "DFLASH speculative decoding only supports CUDA, NPU and CPU devices." ) + _disable_overlap_schedule_for_cpu(server_args) + if resolved_view(server_args).enable_dp_attention: raise ValueError( "Currently DFLASH speculative decoding does not support dp attention." @@ -573,9 +575,14 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None: "triton", "trtllm_mha", "ascend", + "intel_amx", ) - # Use triton on ROCm (no FlashInfer), flashinfer on CUDA. - fallback_backend = "triton" if is_hip() else "flashinfer" + # Use triton on ROCm (no FlashInfer), flashinfer on CUDA, and the AMX + # kernels on CPU. + if cfg.device == "cpu": + fallback_backend = "intel_amx" + else: + fallback_backend = "triton" if is_hip() else "flashinfer" draft_backend = cfg.speculative_draft_attention_backend if draft_backend is None: diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py index 8dc211d1b2a5..9b6976ac4813 100644 --- a/python/sglang/srt/layers/attention/intel_amx_backend.py +++ b/python/sglang/srt/layers/attention/intel_amx_backend.py @@ -5,6 +5,7 @@ import torch from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -61,6 +62,10 @@ def __init__(self, model_runner: ModelRunner): self._attn_logits_buffers: dict[tuple[int, int], torch.Tensor] = {} + # All-visible qlen masks for non-causal (ENCODER_ONLY) spec blocks, + # keyed by (batch_size, draft_token_num). + self._non_causal_masks: dict[tuple[int, int], torch.Tensor] = {} + # speculative decoding params self.num_draft_tokens = get_spec().speculative_num_draft_tokens @@ -108,7 +113,7 @@ def _build_extend_metadata(self, forward_batch: ForwardBatch): # exactly the kernel's built-in causal masking, and skipping the explicit # mask lets extend_attention_cpu take its faster mask-free path. EAGLE # has tree_topk == topk (> 1 for real trees); NGRAM has tree_topk == -1 - # (irregular tree); both need the mask. + # (irregular tree); both need the mask. DFLASH is a chain (== 1). if spec_info.tree_topk != 1: custom_mask = spec_info.custom_mask if custom_mask is not None and custom_mask.numel() > 0: @@ -214,6 +219,19 @@ def forward_extend( # verify batches carry no extend_* fields; see _build_extend_metadata). seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata + if ( + tree_mask is None + and not layer.is_cross_attention + and layer.attn_type == AttentionType.ENCODER_ONLY + and forward_batch.forward_mode.is_target_verify() + ): + # The kernel's implicit mask is causal. A non-causal layer (DFLASH + # draft blocks) needs every query in the block to see every key in + # it, which an all-visible qlen mask expresses. + tree_mask = self._get_non_causal_qlen_mask( + forward_batch.batch_size, forward_batch.spec_info.draft_token_num + ) + _, max_extend_len = self.forward_metadata if seq_lens.dtype != torch.int64: seq_lens = seq_lens.to(torch.int64) @@ -303,6 +321,16 @@ def forward_decode( ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + def _get_non_causal_qlen_mask(self, bs: int, draft_token_num: int) -> torch.Tensor: + key = (int(bs), int(draft_token_num)) + mask = self._non_causal_masks.get(key) + if mask is None: + mask = torch.ones( + key[0] * key[1] * key[1], dtype=torch.bool, device=self.device + ) + self._non_causal_masks[key] = mask + return mask + def _get_attn_logits_buffer( self, num_seqs: int, num_heads: int, v_head_dim: int ) -> torch.Tensor: diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 3a88ad392e1b..d58631b10933 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -39,17 +39,19 @@ from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.speculative.dflash_utils import ( can_dflash_slice_qkv_weight, + dflash_head_logits, get_dflash_attention_sliding_window_size, get_dflash_layer_types, is_dense_head_weight, is_nemotron_35_draft_config, parse_dflash_draft_config, ) -from sglang.srt.utils import is_npu, set_weight_attrs +from sglang.srt.utils import is_cpu, is_npu, set_weight_attrs from sglang.srt.utils.common import get_compiler_backend from sglang.srt.utils.hf_transformers_utils import get_rope_config _is_npu = is_npu() +_is_cpu = is_cpu() if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope logger = logging.getLogger(__name__) @@ -90,8 +92,14 @@ def _project_candidate_logits( ) -> torch.Tensor: """Project draft hiddens through the target head, restricted to the org vocab.""" if not use_quant_head: - weight = lm_head.weight - return torch.matmul(hidden.to(weight.dtype), weight[:num_org].T) + logits = dflash_head_logits(lm_head, hidden, 0, num_org) + if logits.shape[-1] <= num_org: + return logits + # An AMX-prepacked head can only produce the full local shard, so the + # padded tail is masked out of the top-k like the quantized path below. + logits = logits.contiguous() + logits[:, num_org:] = float("-inf") + return logits # A packed weight can't be row-sliced to the org vocab like the dense path, # and flashinfer's radix top-k rejects the crop view (non-contiguous), so # mask the padded tail out of the top-k instead. @@ -219,6 +227,8 @@ def __init__( rotary = self.rotary_emb self.use_table_qk_norm_rope = ( not _is_npu + # table_qk_norm_rope_ is a Triton kernel with no CPU equivalent. + and not _is_cpu and hasattr(rotary, "cos_sin_cache") and getattr(rotary, "rotary_dim", None) == head_dim and getattr(rotary, "is_neox_style", False) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 57a49f6654f6..45d3bdc4b34f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -254,6 +254,7 @@ "triton", "ascend", "trtllm_mha", + "intel_amx", ] add_draft_attention_backend_choices = DRAFT_ATTENTION_BACKEND_CHOICES.extend diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 8390480af01c..9e6a9b8baa8f 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -56,6 +56,12 @@ def __post_init__(self): self.num_tokens_per_req = int(self.draft_token_num) self.num_tokens_for_logprob_per_req = int(self.draft_token_num) + @property + def tree_topk(self) -> int: + # DFLASH proposals are a linear chain, so backends that branch on tree + # width (e.g. intel_amx) can take their mask-free causal path. + return int(self.topk) + def prepare_for_verify( self, batch: ScheduleBatch, diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index db852e486482..2b07e1fb0df8 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -19,6 +19,7 @@ from sglang.srt.managers.schedule_batch import Req from sglang.srt.speculative.spec_utils import sample_simulated_acc_len from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu +from sglang.srt.utils.common import use_intel_amx_backend DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" @@ -34,6 +35,9 @@ "TritonAttnBackend", "TRTLLMHAAttnBackend", "TRTLLMMLABackend", + # The CPU kernel applies causal masking to the verify block on its own + # when the proposal is a linear chain (tree_topk == 1). + "IntelAMXAttnBackend", } ) @@ -735,9 +739,44 @@ def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]: ) if not hasattr(qkv_proj, "weight"): return False, "qkv weight tensor is missing" + if use_intel_amx_backend(qkv_proj): + # The AMX path repacks the weight into a blocked VNNI layout, so its rows + # no longer map to q/k/v output ranges and a plain F.linear against it + # would read the packed bytes as if they were a dense matrix. + return False, "AMX-prepacked qkv_proj weight cannot be sliced" return True, "" +def dflash_head_logits( + lm_head: Any, + hidden_states: torch.Tensor, + start: int = 0, + end: Optional[int] = None, +) -> torch.Tensor: + """Logits for rows ``[start:end)`` of a dense target lm_head weight. + + On CPU with AMX the head weight is prepacked, so the rows cannot be sliced + before the matmul; the packed kernel produces the full local shard and the + requested range is taken from the output instead. + """ + weight = lm_head.weight + if hidden_states.dtype != weight.dtype: + hidden_states = hidden_states.to(weight.dtype) + if use_intel_amx_backend(lm_head): + logits = torch.ops.sgl_kernel.weight_packed_linear( + hidden_states.contiguous(), + weight, + None, # bias + True, # is_vnni + ) + if start != 0 or end is not None: + logits = logits[:, start : logits.shape[-1] if end is None else end] + return logits + if start != 0 or end is not None: + weight = weight[start:end] + return torch.matmul(hidden_states, weight.T) + + def can_dflash_use_fused_qkv_proj(qkv_proj: Any) -> Tuple[bool, str]: """Validate whether a QKV layer is eligible for DFlash fused KV materialization.""" eligible, reason = can_dflash_slice_qkv_weight(qkv_proj) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 7da21952385e..1b2cec2a065c 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -49,6 +49,7 @@ can_dflash_use_fused_qkv_proj, compute_dflash_correct_drafts_and_bonus, compute_dflash_sampling_correct_drafts_and_bonus, + dflash_head_logits, is_dense_head_weight, is_dflash_sampling_verify_available, parse_dflash_draft_config, @@ -70,9 +71,10 @@ assign_req_to_token_pool_func, build_grammar_vocab_mask, ) -from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu +from sglang.srt.utils import get_available_gpu_memory, is_cpu, is_cuda, is_hip, is_npu _is_npu = is_npu() +_is_cpu = is_cpu() logger = logging.getLogger(__name__) @@ -462,7 +464,8 @@ def init_attention_backends(self): def init_cuda_graphs(self): capture_decode_cuda_graph = ( - get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED + not _is_cpu + and get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED ) if is_cuda() and capture_decode_cuda_graph: available_mem = get_available_gpu_memory(self.device, self.gpu_id) @@ -1109,7 +1112,7 @@ def _cast_hs(x: torch.Tensor) -> torch.Tensor: for start in range(0, num_tokens, int(chunk_size)): end = min(num_tokens, start + int(chunk_size)) hs = _cast_hs(hidden_states[start:end]) - logits = torch.matmul(hs, weight.T) + logits = dflash_head_logits(lm_head, hs) out_tokens[start:end] = torch.argmax(logits, dim=-1).to(torch.long) return out_tokens @@ -1164,7 +1167,7 @@ def _ensure_local_reduce_buffers( end = min(num_tokens, start + fast_chunk_size) hs = _cast_hs(hidden_states[start:end]) if num_org > 0: - base_logits = torch.matmul(hs, weight[:num_org].T) + base_logits = dflash_head_logits(lm_head, hs, 0, num_org) local_max, local_arg = _ensure_local_reduce_buffers( end - start, base_logits.dtype, hs.device ) @@ -1182,7 +1185,7 @@ def _ensure_local_reduce_buffers( # Base vocab logits. if num_org > 0: - base_logits = torch.matmul(hs, weight[:num_org].T) + base_logits = dflash_head_logits(lm_head, hs, 0, num_org) local_max, local_arg = _ensure_local_reduce_buffers( chunk_len, base_logits.dtype, hs.device ) @@ -1202,8 +1205,8 @@ def _ensure_local_reduce_buffers( if num_added > 0: added_slice_start = num_org_padded added_slice_end = num_org_padded + num_added - added_logits = torch.matmul( - hs, weight[added_slice_start:added_slice_end].T + added_logits = dflash_head_logits( + lm_head, hs, added_slice_start, added_slice_end ) added_max, added_arg = torch.max(added_logits, dim=-1) use_added = added_max > local_max @@ -1754,10 +1757,11 @@ def forward_batch_generation( ) # `seq_lens` is carried over from the previous overlap iteration and may have been - # produced on another stream. - batch.seq_lens.record_stream( - torch.get_device_module(self.device).current_stream() - ) + # produced on another stream. CPU tensors have no stream to record against. + if not _is_cpu: + batch.seq_lens.record_stream( + torch.get_device_module(self.device).current_stream() + ) bs = len(batch.seq_lens) device = self.device diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index eb9ddb80631d..0b884d881c5a 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -9,7 +9,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode -from sglang.srt.runtime_context import attention_backends, get_spec +from sglang.srt.runtime_context import attention_backends, get_device, get_spec from sglang.srt.server_args import DRAFT_ATTENTION_BACKEND_CHOICES, ServerArgs from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -29,6 +29,13 @@ class DraftWorkerBundle(msgspec.Struct, frozen=True): resolved_attention_backend: str +def _default_draft_attention_backend() -> str: + # A GPU default on a CPU run would fail at the first draft forward. + if get_device().device == "cpu": + return "intel_amx" + return "triton" if torch.version.hip else "flashinfer" + + def _resolve_draft_attention_backend_fallback(*, algo_label: str) -> str: """The draft's attention backend, from the published leaves. @@ -40,9 +47,9 @@ def _resolve_draft_attention_backend_fallback(*, algo_label: str) -> str: if draft_backend is None: draft_backend, _ = attention_backends() if draft_backend is None: - return "triton" if torch.version.hip else "flashinfer" + return _default_draft_attention_backend() if draft_backend not in DRAFT_ATTENTION_BACKEND_CHOICES: - fallback = "triton" if torch.version.hip else "flashinfer" + fallback = _default_draft_attention_backend() logger.warning( "%s draft worker only supports attention_backend in %s for now, " "but got %r. Falling back to '%s'.", diff --git a/test/registered/unit/layers/attention/test_intel_amx_non_causal_mask.py b/test/registered/unit/layers/attention/test_intel_amx_non_causal_mask.py new file mode 100644 index 000000000000..c67ae46f4fff --- /dev/null +++ b/test/registered/unit/layers/attention/test_intel_amx_non_causal_mask.py @@ -0,0 +1,115 @@ +"""The mask intel_amx hands to extend_attention_cpu in TARGET_VERIFY. + +`extend_attention_cpu`'s implicit mask is causal, which is what a decoder layer +wants and what lets the kernel skip the explicit-mask path entirely. A DFLASH +draft block is not causal: its queries are one block that must all see each +other, and that is only expressible as an explicit all-visible qlen mask. +""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.attention.intel_amx_backend import IntelAMXAttnBackend +from sglang.srt.layers.radix_attention import AttentionType +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + +BATCH_SIZE = 2 +DRAFT_TOKEN_NUM = 4 +COMMITTED_LEN = 7 +EXTEND_LEN = 3 +HEAD_NUM = 1 +HEAD_DIM = 8 + + +class _StubKVPool: + def __init__(self): + self.buffer = torch.zeros(16, HEAD_NUM, HEAD_DIM) + + def set_kv_buffer(self, layer, loc, k, v): + pass + + def get_key_buffer(self, layer_id): + return self.buffer + + def get_value_buffer(self, layer_id): + return self.buffer + + +def _layer(attn_type): + return SimpleNamespace( + layer_id=0, + tp_q_head_num=HEAD_NUM, + qk_head_dim=HEAD_DIM, + v_head_dim=HEAD_DIM, + is_cross_attention=False, + attn_type=attn_type, + scaling=1.0, + logit_cap=0.0, + sliding_window_size=-1, + ) + + +def _verify_forward_batch(num_tokens): + return SimpleNamespace( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=BATCH_SIZE, + seq_lens=torch.full((BATCH_SIZE,), COMMITTED_LEN, dtype=torch.int32), + spec_info=SimpleNamespace( + draft_token_num=DRAFT_TOKEN_NUM, tree_topk=1, custom_mask=None + ), + extend_seq_lens=torch.full((BATCH_SIZE,), EXTEND_LEN, dtype=torch.int32), + extend_start_loc=torch.tensor([0, EXTEND_LEN], dtype=torch.int32), + req_pool_indices=torch.arange(BATCH_SIZE, dtype=torch.int32), + out_cache_loc=torch.arange(num_tokens, dtype=torch.int32), + encoder_out_cache_loc=None, + encoder_lens=None, + ) + + +def _tree_mask_reaching_the_kernel(attn_type): + """Drive forward_extend the way a forward pass does and return the mask + argument that reached extend_attention_fwd.""" + num_tokens = BATCH_SIZE * DRAFT_TOKEN_NUM + forward_batch = _verify_forward_batch(num_tokens) + + backend = object.__new__(IntelAMXAttnBackend) + backend.device = "cpu" + backend.token_to_kv_pool = _StubKVPool() + backend.req_to_token_pool = SimpleNamespace( + req_to_token=torch.zeros(BATCH_SIZE, 64, dtype=torch.int32) + ) + backend.swa_out_cache_loc = None + backend.forward_metadata = (None, DRAFT_TOKEN_NUM) + backend.extend_metadata = backend._build_extend_metadata(forward_batch) + backend._non_causal_masks = {} + + seen = {} + backend.extend_attention_fwd = lambda *args: seen.update(tree_mask=args[-1]) + + qkv = torch.zeros(num_tokens, HEAD_NUM * HEAD_DIM) + backend.forward_extend(qkv, qkv, qkv, _layer(attn_type), forward_batch) + return seen["tree_mask"] + + +class TestIntelAMXNonCausalMask(unittest.TestCase): + def test_non_causal_verify_layer_gets_an_all_visible_mask(self): + tree_mask = _tree_mask_reaching_the_kernel(AttentionType.ENCODER_ONLY) + + self.assertEqual(tree_mask.dtype, torch.bool) + self.assertEqual( + tree_mask.shape, (BATCH_SIZE * DRAFT_TOKEN_NUM * DRAFT_TOKEN_NUM,) + ) + self.assertTrue(bool(tree_mask.all())) + + def test_causal_verify_layer_keeps_the_mask_free_path(self): + # Supplying a mask here would cost the kernel's fast path for no gain. + self.assertIsNone(_tree_mask_reaching_the_kernel(AttentionType.DECODER)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_dflash_cpu.py b/test/registered/unit/spec/test_dflash_cpu.py new file mode 100644 index 000000000000..ecd00eeb7300 --- /dev/null +++ b/test/registered/unit/spec/test_dflash_cpu.py @@ -0,0 +1,137 @@ +"""Unit tests for running DFLASH speculative decoding on CPU. + +Covers the pieces that used to be CUDA/NPU-only: the server-arg device gate and +draft attention backend resolution, the linear-chain `tree_topk` that the +intel_amx backend branches on, and the AMX-prepacked weight guards that keep +DFlash from reading packed weights as if they were dense matrices. +""" + +import sys +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.arg_groups.overrides import resolution_result +from sglang.srt.arg_groups.speculative_hook import ( + _disable_overlap_schedule_for_cpu, + _handle_dflash, + _resolve_dflash_draft_attention_backend, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.dflash_info import DFlashVerifyInput +from sglang.srt.speculative.dflash_utils import ( + can_dflash_slice_qkv_weight, + can_dflash_use_fused_qkv_proj, + dflash_head_logits, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _server_args(**overrides) -> ServerArgs: + # ServerArgs(model_path="dummy") early-returns the resolution pipeline, so + # the raw fields below are still writable. + server_args = ServerArgs(model_path="dummy") + server_args.device = "cpu" + server_args.speculative_algorithm = "DFLASH" + server_args.speculative_draft_model_path = "/tmp/draft" + server_args.speculative_draft_attention_backend = None + for key, value in overrides.items(): + setattr(server_args, key, value) + return server_args + + +def test_dflash_device_gate_rejects_an_unsupported_device(): + with pytest.raises(ValueError, match="CUDA, NPU and CPU"): + _handle_dflash(_server_args(device="xpu")) + + +@pytest.mark.parametrize("device", ["cpu", "cuda", "cuda:0", "npu"]) +def test_dflash_device_gate_admits_the_supported_devices(device): + # dp attention is rejected immediately after the device gate, so reaching + # that error is what proves the device itself was accepted. + args = _server_args(device=device, enable_dp_attention=True) + with pytest.raises(ValueError, match="dp attention"): + _handle_dflash(args) + + +def test_dflash_cpu_disables_overlap_schedule(): + args = _server_args() + _disable_overlap_schedule_for_cpu(args) + assert resolution_result(args, "disable_overlap_schedule") is True + + +def test_dflash_leaves_overlap_schedule_alone_off_cpu(): + args = _server_args(device="cuda") + _disable_overlap_schedule_for_cpu(args) + assert resolution_result(args, "disable_overlap_schedule") is not True + + +def test_dflash_cpu_draft_backend_defaults_to_intel_amx(): + # Without this the draft inherits a GPU default and fails at its first + # forward. + args = _server_args() + _resolve_dflash_draft_attention_backend(args) + assert resolution_result(args, "speculative_draft_attention_backend") == "intel_amx" + + +def test_dflash_cpu_draft_backend_is_kept_when_explicit(): + args = _server_args(speculative_draft_attention_backend="intel_amx") + _resolve_dflash_draft_attention_backend(args) + assert resolution_result(args, "speculative_draft_attention_backend") == "intel_amx" + + +def test_dflash_verify_input_is_a_linear_chain(): + # intel_amx skips the explicit mask (and takes the faster kernel path) only + # when tree_topk == 1. + verify_input = DFlashVerifyInput( + draft_token=torch.zeros(4, dtype=torch.long), + positions=torch.zeros(4, dtype=torch.long), + draft_token_num=4, + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + assert verify_input.tree_topk == 1 + + +def _fake_qkv_proj(*, packed: bool): + proj = SimpleNamespace( + quant_method=UnquantizedLinearMethod(), + weight=torch.zeros(6, 4), + bias=None, + ) + if packed: + proj.use_intel_amx_backend = True + return proj + + +def test_amx_prepacked_qkv_cannot_be_sliced(): + eligible, reason = can_dflash_slice_qkv_weight(_fake_qkv_proj(packed=True)) + assert not eligible + assert "AMX" in reason + assert not can_dflash_use_fused_qkv_proj(_fake_qkv_proj(packed=True))[0] + + assert can_dflash_slice_qkv_weight(_fake_qkv_proj(packed=False))[0] + + +def test_dflash_head_logits_dense_path_matches_matmul(): + hidden = torch.randn(3, 8) + lm_head = SimpleNamespace(weight=torch.randn(10, 8)) + + torch.testing.assert_close( + dflash_head_logits(lm_head, hidden), hidden @ lm_head.weight.T + ) + torch.testing.assert_close( + dflash_head_logits(lm_head, hidden, 0, 6), hidden @ lm_head.weight[:6].T + ) + torch.testing.assert_close( + dflash_head_logits(lm_head, hidden, 6, 10), hidden @ lm_head.weight[6:10].T + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))