@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 54254f7a61

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 22:45 UTC

[ROCm][CI] Fix spec decode logprobs flakiness and parametrize tree attention backends (#34599)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Feb 20, 2026 at 22:25 UTC 54254f7a6155002b1a493c3cefda9752bc9ce92f

2 files changed +356 -151

tests/v1/sample/test_logprobs.py

+158 -134

@@ -52,7 +52,7 @@ def vllm_model(vllm_runner, request) -> Generator[VllmRunner, None, None]:

52 # TODO: enable this once we support it for

53 # prompt logprobs.

54 enable_prefix_caching=request.param,

55 - gpu_memory_utilization=0.4, # up to 2 alive concurrently

55 + gpu_memory_utilization=0.4,

56 ) as vllm_model:

57 yield vllm_model

@@ -366,21 +366,20 @@ def test_max_logprobs():

366 Should also fail for `prompt_logprobs > max_logprobs`

367 APC should not matter as this test checks basic request validation.

368 """

369 - runner = VllmRunner(

369 + with VllmRunner(

370 "facebook/opt-125m",

371 max_logprobs=1,

372 enable_prefix_caching=False,

373 - # 2 other llms alive during whole session

373 gpu_memory_utilization=0.15,

374 max_model_len=256,

376 - )

377 - vllm_sampling_params = SamplingParams(logprobs=1)

378 - # should pass

379 - runner.generate(["Hello world"], sampling_params=vllm_sampling_params)

375 + ) as runner:

376 + vllm_sampling_params = SamplingParams(logprobs=1)

377 + # should pass

378 + runner.generate(["Hello world"], sampling_params=vllm_sampling_params)

379

381 - bad_sampling_params = SamplingParams(logprobs=2)

382 - with pytest.raises(ValueError):

383 - runner.generate(["Hello world"], sampling_params=bad_sampling_params)

380 + bad_sampling_params = SamplingParams(logprobs=2)

381 + with pytest.raises(ValueError):

382 + runner.generate(["Hello world"], sampling_params=bad_sampling_params)

383

384

385 def test_none_logprobs(vllm_model, example_prompts):

@@ -449,33 +448,31 @@ def test_all_logprobs(example_prompts):

448 Args:

449 example_prompts: list of example prompts (test fixture)

450 """

452 - runner = VllmRunner(

451 + with VllmRunner(

452 "facebook/opt-125m",

453 max_logprobs=-1,

454 enable_prefix_caching=False,

456 - # 2 other llms alive during whole session

455 gpu_memory_utilization=0.15,

456 max_model_len=256,

459 - )

460 -

461 - sampling_params_logprobs_all = SamplingParams(

462 - max_tokens=5, logprobs=-1, prompt_logprobs=-1

463 - )

464 - results_logprobs_all = runner.llm.generate(

465 - example_prompts, sampling_params=sampling_params_logprobs_all

466 - )

467 - vocab_size = runner.llm.llm_engine.model_config.get_vocab_size()

457 + ) as runner:

458 + sampling_params_logprobs_all = SamplingParams(

459 + max_tokens=5, logprobs=-1, prompt_logprobs=-1

460 + )

461 + results_logprobs_all = runner.llm.generate(

462 + example_prompts, sampling_params=sampling_params_logprobs_all

463 + )

464 + vocab_size = runner.llm.llm_engine.model_config.get_vocab_size()

465

469 - for i in range(len(results_logprobs_all)):

470 - logprobs = results_logprobs_all[i].outputs[0].logprobs

471 - prompt_logprobs = results_logprobs_all[i].prompt_logprobs

472 - assert logprobs is not None

473 - for logprob in logprobs:

474 - assert len(logprob) == vocab_size

475 - assert prompt_logprobs is not None

476 - assert prompt_logprobs[0] is None

477 - for prompt_logprob in prompt_logprobs[1:]:

478 - assert len(prompt_logprob) == vocab_size

466 + for i in range(len(results_logprobs_all)):

467 + logprobs = results_logprobs_all[i].outputs[0].logprobs

468 + prompt_logprobs = results_logprobs_all[i].prompt_logprobs

469 + assert logprobs is not None

470 + for logprob in logprobs:

471 + assert len(logprob) == vocab_size

472 + assert prompt_logprobs is not None

473 + assert prompt_logprobs[0] is None

474 + for prompt_logprob in prompt_logprobs[1:]:

475 + assert len(prompt_logprob) == vocab_size

476

477

478 @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))

@@ -495,24 +492,28 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

492 max_model_len=16,

493 logprobs_mode=logprobs_mode,

494 )

498 - vllm_sampling_params = SamplingParams(logprobs=1)

499 - results = llm.generate(["Hello world"], sampling_params=vllm_sampling_params)

500 -

501 - total_token_with_logprobs = 0

502 - positive_values = 0

503 - for output in results[0].outputs:

504 - for logprobs in output.logprobs:

505 - for token_id in logprobs:

506 - logprob = logprobs[token_id]

507 - if logprobs_mode in ("raw_logprobs", "processed_logprobs"):

508 - assert logprob.logprob <= 0

509 - if logprob.logprob > 0:

510 - positive_values = positive_values + 1

511 - total_token_with_logprobs = total_token_with_logprobs + 1

512 - assert total_token_with_logprobs >= len(results[0].outputs)

513 - if logprobs_mode in ("raw_logits", "processed_logits"):

514 - assert positive_values > 0

515 - del llm

495 + try:

496 + vllm_sampling_params = SamplingParams(logprobs=1)

497 + results = llm.generate(["Hello world"], sampling_params=vllm_sampling_params)

498 +

499 + total_token_with_logprobs = 0

500 + positive_values = 0

501 + for output in results[0].outputs:

502 + for logprobs in output.logprobs:

503 + for token_id in logprobs:

504 + logprob = logprobs[token_id]

505 + if logprobs_mode in ("raw_logprobs", "processed_logprobs"):

506 + assert logprob.logprob <= 0

507 + if logprob.logprob > 0:

508 + positive_values = positive_values + 1

509 + total_token_with_logprobs = total_token_with_logprobs + 1

510 + assert total_token_with_logprobs >= len(results[0].outputs)

511 + if logprobs_mode in ("raw_logits", "processed_logits"):

512 + assert positive_values > 0

513 + finally:

514 + del llm

515 + torch.cuda.empty_cache()

516 + cleanup_dist_env_and_memory()

517

518

519 class TestCorrectDecodedToken:

@@ -767,7 +768,7 @@ class TestCorrectDecodedToken:

768 # Simulate cases where individual tokens decode to "�"

769 # but combinations decode correctly

770 if len(ids) == 1:

770 - if ids[0] == 3 or ids[0] == 4 or ids[0] == 8 or ids[0] == 9:

771 + if ids[0] in (3, 4, 8, 9):

772 return "�"

773 elif len(ids) == 2:

774 if ids == [2, 3]:

@@ -809,42 +810,41 @@ def test_verify_tokens_integration():

810 corrects tokens ending with the replacement character "�".

811 Uses facebook/opt-125m which is known to produce these issues.

812 """

812 - runner = VllmRunner(

813 + with VllmRunner(

814 "facebook/opt-125m",

815 max_logprobs=0,

816 enable_prefix_caching=False,

817 gpu_memory_utilization=0.15,

818 max_model_len=256,

818 - )

819 -

820 - # Use a prompt that triggers multi-byte UTF-8 issues

821 - # Based on user's example: "In this example,"

822 - test_prompts = ["In this example,"]

823 -

824 - sampling_params = SamplingParams(

825 - max_tokens=16,

826 - temperature=0,

827 - logprobs=0,

828 - )

819 + ) as runner:

820 + # Use a prompt that triggers multi-byte UTF-8 issues

821 + # Based on user's example: "In this example,"

822 + test_prompts = ["In this example,"]

823 +

824 + sampling_params = SamplingParams(

825 + max_tokens=16,

826 + temperature=0,

827 + logprobs=0,

828 + )

829

830 - results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

831 -

832 - # Verify that decoded tokens don't contain replacement characters

833 - for result in results:

834 - assert result.outputs[0].logprobs is not None

835 - for logprob_dict in result.outputs[0].logprobs:

836 - for token_id, logprob_info in logprob_dict.items():

837 - decoded_token = logprob_info.decoded_token

838 - # Decoded tokens should not end with replacement character

839 - # They should either be corrected or empty string

840 - assert not decoded_token.endswith("�"), (

841 - f"Token {token_id} decoded to '{decoded_token}' which "

842 - f"ends with replacement character"

843 - )

844 - # Decoded tokens should not contain lone replacement characters

845 - assert decoded_token != "�", (

846 - f"Token {token_id} is a lone replacement character"

847 - )

830 + results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

831 +

832 + # Verify that decoded tokens don't contain replacement characters

833 + for result in results:

834 + assert result.outputs[0].logprobs is not None

835 + for logprob_dict in result.outputs[0].logprobs:

836 + for token_id, logprob_info in logprob_dict.items():

837 + decoded_token = logprob_info.decoded_token

838 + # Decoded tokens should not end with replacement character

839 + # They should either be corrected or empty string

840 + assert not decoded_token.endswith("�"), (

841 + f"Token {token_id} decoded to '{decoded_token}' which "

842 + f"ends with replacement character"

843 + )

844 + # Decoded tokens should not contain lone replacement characters

845 + assert decoded_token != "�", (

846 + f"Token {token_id} is a lone replacement character"

847 + )

848

849

850 def test_utf8_edge_cases_with_real_model():

@@ -853,45 +853,44 @@ def test_utf8_edge_cases_with_real_model():

853 Tests prompts that are likely to trigger byte-fallback tokenization

854 and multi-byte UTF-8 splitting.

855 """

856 - runner = VllmRunner(

856 + with VllmRunner(

857 "facebook/opt-125m",

858 max_logprobs=1,

859 enable_prefix_caching=False,

860 gpu_memory_utilization=0.15,

861 max_model_len=256,

862 - )

863 -

864 - # Prompts with various multi-byte UTF-8 characters

865 - test_prompts = [

866 - 'Smart quotes: "Hello"', # Curly quotes

867 - "Em dash — test", # Em dash

868 - "Ellipsis… continues", # Ellipsis

869 - "Chinese: 你好", # Chinese characters

870 - "Emoji: 😀 🎉", # Emojis

871 - 'Mixed: "quoted" — with symbols', # Mixed

872 - ]

873 -

874 - sampling_params = SamplingParams(

875 - max_tokens=10,

876 - temperature=0,

877 - logprobs=1,

878 - )

862 + ) as runner:

863 + # Prompts with various multi-byte UTF-8 characters

864 + test_prompts = [

865 + 'Smart quotes: "Hello"', # Curly quotes

866 + "Em dash — test", # Em dash

867 + "Ellipsis… continues", # Ellipsis

868 + "Chinese: 你好", # Chinese characters

869 + "Emoji: 😀 🎉", # Emojis

870 + 'Mixed: "quoted" — with symbols', # Mixed

871 + ]

872 +

873 + sampling_params = SamplingParams(

874 + max_tokens=10,

875 + temperature=0,

876 + logprobs=1,

877 + )

878

880 - results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

879 + results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

880

882 - for i, result in enumerate(results):

883 - prompt = test_prompts[i]

884 - assert result.outputs[0].logprobs is not None

881 + for i, result in enumerate(results):

882 + prompt = test_prompts[i]

883 + assert result.outputs[0].logprobs is not None

884

886 - # Check that no decoded tokens end with replacement character

887 - for logprob_dict in result.outputs[0].logprobs:

888 - for token_id, logprob_info in logprob_dict.items():

889 - decoded_token = logprob_info.decoded_token

890 - assert not decoded_token.endswith("�"), (

891 - f"Prompt: '{prompt}'\n"

892 - f"Token {token_id} decoded to '{decoded_token}' which "

893 - f"ends with replacement character"

894 - )

885 + # Check that no decoded tokens end with replacement character

886 + for logprob_dict in result.outputs[0].logprobs:

887 + for token_id, logprob_info in logprob_dict.items():

888 + decoded_token = logprob_info.decoded_token

889 + assert not decoded_token.endswith("�"), (

890 + f"Prompt: '{prompt}'\n"

891 + f"Token {token_id} decoded to '{decoded_token}' which "

892 + f"ends with replacement character"

893 + )

894

895

896 def test_correct_decoded_token_preserves_valid_tokens():

@@ -901,36 +900,35 @@ def test_correct_decoded_token_preserves_valid_tokens():

900 ending with "�", but this test verifies the broader _verify_tokens

901 logic doesn't affect valid tokens.

902 """

904 - runner = VllmRunner(

903 + with VllmRunner(

904 "facebook/opt-125m",

905 max_logprobs=2,

906 enable_prefix_caching=False,

907 gpu_memory_utilization=0.15,

908 max_model_len=256,

910 - )

911 -

912 - # Simple prompt with standard ASCII characters

913 - test_prompts = ["Hello world, this is a test."]

914 -

915 - sampling_params = SamplingParams(

916 - max_tokens=10,

917 - temperature=0,

918 - logprobs=2,

919 - )

909 + ) as runner:

910 + # Simple prompt with standard ASCII characters

911 + test_prompts = ["Hello world, this is a test."]

912 +

913 + sampling_params = SamplingParams(

914 + max_tokens=10,

915 + temperature=0,

916 + logprobs=2,

917 + )

918

921 - results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

919 + results = runner.llm.generate(test_prompts, sampling_params=sampling_params)

920

923 - for result in results:

924 - assert result.outputs[0].logprobs is not None

921 + for result in results:

922 + assert result.outputs[0].logprobs is not None

923

926 - # All decoded tokens should be valid strings

927 - for logprob_dict in result.outputs[0].logprobs:

928 - for token_id, logprob_info in logprob_dict.items():

929 - decoded_token = logprob_info.decoded_token

930 - # Valid tokens should be non-empty strings (or empty if corrected)

931 - assert isinstance(decoded_token, str)

932 - # Should not contain replacement character

933 - assert "�" not in decoded_token

924 + # All decoded tokens should be valid strings

925 + for logprob_dict in result.outputs[0].logprobs:

926 + for token_id, logprob_info in logprob_dict.items():

927 + decoded_token = logprob_info.decoded_token

928 + # Valid tokens should be non-empty strings (or empty if corrected)

929 + assert isinstance(decoded_token, str)

930 + # Should not contain replacement character

931 + assert "�" not in decoded_token

932

933

934 @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))

@@ -985,16 +983,33 @@ def test_correct_decoded_token_preserves_valid_tokens():

983 def test_spec_decode_logprobs(

984 logprobs_mode: LogprobsMode,

985 model_setup: tuple[str, str, dict, int],

986 + monkeypatch,

987 ):

988 """Spec decode logprobs should match those of the base model.

989

990 + Runs the base model and spec decode model sequentially, ensuring

991 + only one LLM instance is alive at a time to avoid GPU memory

992 + contention. Both use identical chunked prefill settings and eager

993 + mode to control for infrastructure differences.

994 +

995 Args:

996 logprobs_mode: logprobs mode.

997 model_setup: Tuple of (method, base model name,

998 speculative_config dict, top_logprobs).

999 + monkeypatch: pytest fixture for setting env vars.

1000 """

1001 from vllm import LLM

1002

1003 + # The ROCm skinny GEMM kernels (gemm_kernels.cu) are

1004 + # non-deterministic across LLM instantiations due to persistent

1005 + # workgroup scheduling and wave-level shuffle reductions, which

1006 + # causes logprob differences that get misattributed to spec decode.

1007 + # Disable them so this test isolates spec decode correctness only.

1008 + # TODO(akaratza): Remove this workaround once the follow-up to

1009 + # https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975

1010 + # lands with a determinism fix for wvSplitK kernels.

1011 + monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")

1012 +

1013 method, model_name, spec_config, top_logprobs = model_setup

1014

1015 prompt = "Hello world " * 50

@@ -1068,8 +1083,17 @@ def test_spec_decode_logprobs(

1083 for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs):

1084 assert math.isclose(

1085 ref_logprob.logprob, spec_logprob.logprob, rel_tol=5e-2, abs_tol=1e-1

1086 + ), (

1087 + f"Logprob mismatch: ref={ref_logprob.logprob} "

1088 + f"spec={spec_logprob.logprob} "

1089 + f"diff={abs(ref_logprob.logprob - spec_logprob.logprob)} "

1090 + f"(token={ref_logprob.decoded_token!r})"

1091 + )

1092 + assert ref_logprob.rank == spec_logprob.rank, (

1093 + f"Rank mismatch: ref={ref_logprob.rank} "

1094 + f"spec={spec_logprob.rank} "

1095 + f"(token={ref_logprob.decoded_token!r})"

1096 )

1072 - assert ref_logprob.rank == spec_logprob.rank

1097 assert ref_logprob.decoded_token == spec_logprob.decoded_token

1098

1099

tests/v1/spec_decode/test_tree_attention.py

+198 -17

@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (

13 try_get_attention_backend,

14 )

15 from vllm.config import ParallelConfig, SpeculativeConfig

16 +from vllm.platforms import current_platform

17 from vllm.v1.attention.backend import CommonAttentionMetadata

18 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available

19 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -23,11 +24,156 @@ if not is_flash_attn_varlen_func_available():

24 allow_module_level=True,

25 )

27 +# --------------------------------------------------------------------------- #

28 +# KV cache layout adaptation

29 +# --------------------------------------------------------------------------- #

30 +# Two KV cache layouts exist across backends:

31 +#

32 +# Flash layout: (2, num_blocks, block_size, num_kv_heads, head_size)

33 +# - dim 0 separates key (index 0) and value (index 1)

34 +# - Used by: FLASH_ATTN, TREE_ATTN, ROCM_AITER_FA, ROCM_ATTN

35 +#

36 +# Block layout: (num_blocks, 2, block_size, num_kv_heads, head_size)

37 +# - dim 1 separates key (index 0) and value (index 1)

38 +# - Used by: TRITON_ATTN

39 +#

40 +# The test creates KV caches in flash layout (the canonical format used by

41 +# tree attention). When a reference backend needs block layout we transpose

42 +# dims 0 and 1.

43 +#

44 +# Note: ROCM_ATTN uses flash layout for storage but its forward path calls

45 +# PagedAttention.split_kv_cache which reinterprets the raw memory as paged

46 +# layout (num_blocks, num_kv_heads, head_size//x, block_size, x). This is

47 +# a view-level incompatibility, not a transpose - see the TODO in

48 +# _get_available_reference_backends for details.

49 +#

50 +# TODO: Replace this mapping with a `KV_CACHE_LAYOUT` class attribute on each

51 +# AttentionImpl so the layout is self-documented by the backend itself, e.g.:

52 +# class TritonAttentionImpl(AttentionImpl):

53 +# KV_CACHE_LAYOUT = "block"

54 +# --------------------------------------------------------------------------- #

55 +

56 +_BLOCK_KV_LAYOUT_BACKENDS = frozenset(

57 + {

58 + AttentionBackendEnum.TRITON_ATTN,

59 + }

60 +)

61 +

62 +# Backends whose do_kv_cache_update requires engine-level state (e.g.

63 +# ForwardContext) that is not available in this test harness, but whose

64 +# KV cache is flash layout and can be written with reshape_and_cache_flash.

65 +# When a backend is listed here, forward_attention() bypasses

66 +# do_kv_cache_update and writes directly to the cache.

67 +_NEEDS_DIRECT_CACHE_UPDATE = frozenset(

68 + {

69 + AttentionBackendEnum.ROCM_AITER_FA,

70 + }

71 +)

72 +

73 +# Backends with known test-harness incompatibilities - see the TODOs

74 +# inside _get_available_reference_backends for details.

75 +_INCOMPATIBLE_REFERENCE_BACKENDS = frozenset(

76 + {

77 + AttentionBackendEnum.ROCM_AITER_FA,

78 + AttentionBackendEnum.ROCM_ATTN,

79 + }

80 +)

81 +

82 +

83 +def _adapt_kv_cache_for_backend(

84 + kv_cache: torch.Tensor,

85 + backend: AttentionBackendEnum,

86 +) -> torch.Tensor:

87 + """Convert kv_cache from flash layout ``(2, num_blocks, ...)`` to block

88 + layout ``(num_blocks, 2, ...)`` if the backend requires it. Returns the

89 + original tensor unchanged when no conversion is needed."""

90 + if backend in _BLOCK_KV_LAYOUT_BACKENDS:

91 + return kv_cache.transpose(0, 1).contiguous()

92 + return kv_cache

93 +

94 +

95 +def _get_platform_default_backend() -> AttentionBackendEnum:

96 + """Ask the platform what backend it would auto-select at runtime."""

97 + from vllm.v1.attention.selector import AttentionSelectorConfig

98 +

99 + config = AttentionSelectorConfig(

100 + block_size=32,

101 + kv_cache_dtype="auto",

102 + use_mla=False,

103 + use_sparse=False,

104 + head_size=128,

105 + dtype=torch.bfloat16,

106 + )

107 + backend_path = current_platform.get_attn_backend_cls(

108 + selected_backend=None,

109 + attn_selector_config=config,

110 + )

111 + for backend in AttentionBackendEnum:

112 + try:

113 + if backend.get_path() == backend_path:

114 + return backend

115 + except ValueError:

116 + continue

117 + raise RuntimeError(

118 + f"Platform returned backend path '{backend_path}' "

119 + f"that doesn't match any AttentionBackendEnum member."

120 + )

121 +

122 +

123 +def _get_available_reference_backends() -> list[AttentionBackendEnum]:

124 + """Collect all reference backends the current platform can run.

125 +

126 + On CUDA this is just FLASH_ATTN. On ROCm this includes the platform

127 + default plus every backend the hardware supports, so the test validates

128 + tree attention against all of them.

129 + """

130 + if current_platform.is_rocm():

131 + backends: list[AttentionBackendEnum] = []

132 +

133 + # 1. Whatever the platform would auto-select at runtime.

134 + default_backend = _get_platform_default_backend()

135 + if default_backend not in _INCOMPATIBLE_REFERENCE_BACKENDS:

136 + backends.append(default_backend)

137 +

138 + # 2. TRITON_ATTN - always available on ROCm.

139 + if AttentionBackendEnum.TRITON_ATTN not in backends:

140 + backends.append(AttentionBackendEnum.TRITON_ATTN)

141 +

142 + # TODO: Enable ROCM_ATTN. Its forward path uses

143 + # PagedAttention.split_kv_cache which reinterprets the raw

144 + # cache memory as paged layout:

145 + # key: (num_blocks, num_kv_heads, head_size//x, block_size, x)

146 + # value: (num_blocks, num_kv_heads, head_size, block_size)

147 + # Tree attention writes prefix data in NHD flash layout, so the

148 + # same bytes produce completely different values when read in

149 + # paged format. Supporting ROCM_ATTN would require writing

150 + # prefix data via PagedAttention.write_to_paged_cache into a

151 + # separate paged-format KV cache.

152 +

153 + # TODO: Enable ROCM_AITER_FA. Its metadata builder reads head

154 + # counts from the model config at construction time and

155 + # allocates extend_workspace with those dimensions. The test

156 + # uses independent head count parameters (num_heads=2/4,

157 + # num_kv_heads=2) that don't match the model config

158 + # (Llama-3-8B: 32 q heads, 8 kv heads), causing a head count

159 + # mismatch in flash_attn_varlen_func during extend_forward.

160 + # Fixing this requires either matching test head counts to the

161 + # model config or decoupling the builder from model config

162 + # head geometry. The direct cache update path

163 + # (_NEEDS_DIRECT_CACHE_UPDATE) is already in place for when

164 + # this is resolved.

165 +

166 + return backends

167 +

168 + # CUDA: flash attention.

169 + return [AttentionBackendEnum.FLASH_ATTN]

170 +

171

172 class MockAttentionLayer(torch.nn.Module):

173 _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

174 _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

175 _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

176 + layer_name = "mock_layer"

177

178 def __init__(self):

179 super().__init__()

@@ -48,6 +194,13 @@ def forward_attention(

194 spec_token_tree: str | None = None,

195 num_spec_tokens: int = 0,

196 ) -> torch.Tensor:

197 + """Run a single attention forward pass through the given backend.

198 +

199 + ``kv_cache`` is expected in **flash layout**

200 + ``(2, num_blocks, block_size, num_kv_heads, head_size)``.

201 + It is automatically converted when the target backend needs a

202 + different layout.

203 + """

204 batch_size, q_len, num_heads, dim_per_head = q.shape

205 num_kv_heads = k.shape[-2]

206 # Initialize the query and KV sequence lengths.

@@ -116,31 +269,58 @@ def forward_attention(

269 kv_cache_dtype="auto",

270 )

271

272 + # Adapt KV cache layout for this backend.

273 + adapted_kv_cache = _adapt_kv_cache_for_backend(kv_cache, backend)

274 +

275 # Run forward pass and return output.

276 query = q.view(-1, num_heads, dim_per_head)

277 key = k.view(-1, num_kv_heads, dim_per_head)

278 value = v.view(-1, num_kv_heads, dim_per_head)

279 output = torch.empty_like(query)

280 if not try_backend_includes_kv_cache_update(backend):

125 - instance.do_kv_cache_update(

126 - layer=layer,

127 - key=key,

128 - value=value,

129 - kv_cache=kv_cache,

130 - slot_mapping=attn_metadata.slot_mapping,

131 - )

281 + if backend in _NEEDS_DIRECT_CACHE_UPDATE:

282 + # This backend's do_kv_cache_update requires engine-level

283 + # ForwardContext that isn't available in this test harness.

284 + # Write directly using reshape_and_cache_flash since the

285 + # KV cache layout is identical (flash layout, unbind on dim 0).

286 + key_cache, value_cache = adapted_kv_cache.unbind(0)

287 + torch.ops._C_cache_ops.reshape_and_cache_flash(

288 + key,

289 + value,

290 + key_cache,

291 + value_cache,

292 + attn_metadata.slot_mapping,

293 + "auto",

294 + layer._k_scale,

295 + layer._v_scale,

296 + )

297 + else:

298 + instance.do_kv_cache_update(

299 + layer=layer,

300 + key=key,

301 + value=value,

302 + kv_cache=adapted_kv_cache,

303 + slot_mapping=attn_metadata.slot_mapping,

304 + )

305 return instance.forward(

306 layer=layer,

307 query=query,

308 key=key,

309 value=value,

137 - kv_cache=kv_cache.clone(),

310 + kv_cache=adapted_kv_cache.clone(),

311 attn_metadata=attn_metadata,

312 output=output,

313 )

314

315

143 -def test_tree_attn_correctness() -> None:

316 +@pytest.mark.parametrize(

317 + "reference_backend",

318 + _get_available_reference_backends(),

319 + ids=lambda b: b.name,

320 +)

321 +def test_tree_attn_correctness(

322 + reference_backend: AttentionBackendEnum,

323 +) -> None:

324 torch.manual_seed(42)

325 torch.cuda.manual_seed_all(42)

326

@@ -205,7 +385,9 @@ def test_tree_attn_correctness() -> None:

385 dtype=torch.bfloat16,

386 )

387

208 - # Set up the block table and KV cache for paged KV.

388 + # KV cache in flash layout - the canonical format for

389 + # tree attention. forward_attention() handles conversion

390 + # when needed.

391 assert max_sequence_length % block_size == 0

392 max_blocks_per_batch = max_sequence_length // block_size

393 kv_cache = torch.randn(

@@ -263,9 +445,7 @@ def test_tree_attn_correctness() -> None:

445 num_spec_tokens=tree_size_q - 1,

446 ).view(batch_size, -1, num_heads, dim_per_head)

447

266 - # Verify that the chain attention output for each

267 - # branch of the tree (computed using FA3) matches

268 - # the tree attention output.

448 + # Verify each branch against the reference backend.

449 for q_index in range(tree_size_q):

450 # Get the q, k, and v for the branch.

451 branch_mask = tree_attn_mask[q_index, :]

@@ -286,8 +466,8 @@ def test_tree_attn_correctness() -> None:

466 branch_positions, block_table, block_size

467 )

468

289 - # Compute flash attention for the branch.

290 - flash_attn_output = forward_attention(

469 + # Reference attention for this branch.

470 + ref_output = forward_attention(

471 q=q_branch,

472 k=k_branch,

473 v=v_branch,

@@ -295,16 +475,17 @@ def test_tree_attn_correctness() -> None:

475 block_table=block_table,

476 slot_mapping=branch_slot_mapping,

477 seqlen_k=sequence_position + q_len,

298 - backend=AttentionBackendEnum.FLASH_ATTN,

478 + backend=reference_backend,

479 ).view(batch_size, -1, num_heads, dim_per_head)

480

481 # Compare the outputs.

482 assert torch.allclose(

483 tree_attn_output[:, branch_indices],

304 - flash_attn_output,

484 + ref_output,

485 atol=7.81e-3,

486 ), (

487 f"outputs are not close for "

488 + f"reference_backend: {reference_backend.name}, "

489 f"batch_size: {batch_size}, "

490 f"num_heads: {num_heads}, "

491 f"sequence_position: {sequence_position}, "