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

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 6, 07:00 UTC

[Attention] Support distinguishing between short extends and decodes (#37303)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Mar 20, 2026 at 10:49 UTC e1d85e5c2454bd8d349dbe676679380cbe0e920a

9 files changed +178 -135

.buildkite/test_areas/engine.yaml

+12

@@ -70,3 +70,15 @@ steps:

70 device: mi325_4

71 depends_on:

72 - image-build-amd

73 +

74 +- label: V1 e2e (4xH100)

75 + timeout_in_minutes: 60

76 + device: h100

77 + num_devices: 4

78 + optional: true

79 + source_file_dependencies:

80 + - vllm/v1/attention/backends/utils.py

81 + - vllm/v1/worker/gpu_model_runner.py

82 + - tests/v1/e2e/test_hybrid_chunked_prefill.py

83 + commands:

84 + - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py

tests/v1/attention/test_batch_reordering.py

+86 -37

@@ -10,9 +10,10 @@ from vllm.v1.attention.backends.utils import reorder_batch_to_split_decodes_and_

12 class MockInputBatch:

13 - def __init__(self, req_ids, num_computed_tokens_cpu):

13 + def __init__(self, req_ids, num_computed_tokens_cpu, num_prompt_tokens):

14 self.req_ids = req_ids

15 self.num_computed_tokens_cpu = num_computed_tokens_cpu

16 + self.num_prompt_tokens = num_prompt_tokens

18 def swap_states(self, i, j):

19 self.req_ids[i], self.req_ids[j] = self.req_ids[j], self.req_ids[i]

@@ -20,6 +21,10 @@ class MockInputBatch:

21 self.num_computed_tokens_cpu[j],

22 self.num_computed_tokens_cpu[i],

23 )

24 + self.num_prompt_tokens[i], self.num_prompt_tokens[j] = (

25 + self.num_prompt_tokens[j],

26 + self.num_prompt_tokens[i],

27 + )

30 class MockSchedulerOutput:

@@ -29,96 +34,139 @@ class MockSchedulerOutput:

35 @dataclass

36 class ReorderTestCase:

32 - requests: list[tuple[int, int]] # (num_scheduled_tokens, num_computed_tokens)

37 + # (num_scheduled_tokens, num_computed_tokens, num_prompt_tokens)

38 + requests: list[tuple[int, int, int]]

39 expected_order: list[int]

40 expected_modified: bool

41 decode_threshold: int = 1

44 # Test cases for batch reordering

45 +# Format: (num_scheduled, num_computed, num_prompt)

46 REORDER_TEST_CASES = {

47 "all_decodes": ReorderTestCase(

41 - requests=[(1, 10), (1, 20), (1, 30)],

48 + requests=[(1, 10, 10), (1, 20, 20), (1, 30, 30)],

49 expected_order=[0, 1, 2],

50 expected_modified=False,

51 ),

45 - "all_prefills": ReorderTestCase(

46 - requests=[(100, 100), (200, 200), (300, 300)],

52 + "all_long_extends": ReorderTestCase(

53 + requests=[(100, 100, 100), (200, 200, 200), (300, 300, 300)],

54 expected_order=[0, 1, 2],

55 expected_modified=False,

56 ),

50 - "mixed_interleaved": ReorderTestCase(

51 - requests=[(100, 100), (1, 10), (200, 200), (1, 20)],

52 - expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place

57 + "mixed_decodes_long_extends": ReorderTestCase(

58 + requests=[(100, 100, 100), (1, 10, 10), (200, 200, 200), (1, 20, 20)],

59 + expected_order=[3, 1, 2, 0],

60 expected_modified=True,

61 ),

62 "already_ordered": ReorderTestCase(

56 - requests=[(1, 10), (1, 20), (100, 100), (200, 0)],

63 + requests=[(1, 10, 10), (1, 20, 20), (100, 100, 100), (200, 0, 200)],

64 expected_order=[0, 1, 2, 3],

65 expected_modified=False,

66 ),

67 "single_request": ReorderTestCase(

61 - requests=[(1, 10)],

68 + requests=[(1, 10, 10)],

69 expected_order=[0],

70 expected_modified=False,

71 ),

72 "higher_threshold": ReorderTestCase(

66 - requests=[(2, 10), (3, 20), (5, 30), (6, 40)],

73 + requests=[(2, 10, 10), (3, 20, 20), (5, 30, 30), (6, 40, 40)],

74 expected_order=[0, 1, 2, 3],

75 expected_modified=False,

76 decode_threshold=4,

77 ),

78 "decodes_at_end": ReorderTestCase(

72 - requests=[(100, 100), (200, 200), (1, 10), (1, 20)],

79 + requests=[(100, 100, 100), (200, 200, 200), (1, 10, 10), (1, 20, 20)],

80 expected_order=[2, 3, 0, 1],

81 expected_modified=True,

82 ),

76 - "decode_extend_prefill": ReorderTestCase(

77 - requests=[(100, 0), (10, 50), (1, 10)],

83 + "decode_long_extend_prefill": ReorderTestCase(

84 + requests=[(100, 0, 100), (10, 50, 50), (1, 10, 10)],

85 expected_order=[2, 1, 0],

86 expected_modified=True,

87 ),

81 - "extend_prefill_only": ReorderTestCase(

82 - requests=[(100, 0), (10, 50), (200, 0), (20, 75)],

83 - expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place

88 + "long_extend_prefill_only": ReorderTestCase(

89 + requests=[(100, 0, 100), (10, 50, 50), (200, 0, 200), (20, 75, 75)],

90 + expected_order=[3, 1, 2, 0],

91 expected_modified=True,

92 ),

86 - "complicated_mixed_interleaved": ReorderTestCase(

93 + "complicated_mixed": ReorderTestCase(

94 requests=[

88 - (1, 20),

89 - (1, 50),

90 - (374, 0),

91 - (300, 20),

92 - (1, 20),

93 - (256, 0),

94 - (1, 5),

95 - (27, 0),

96 - (1, 4),

95 + (1, 20, 20), # decode

96 + (1, 50, 50), # decode

97 + (374, 0, 374), # prefill

98 + (300, 20, 20), # long_extend

99 + (1, 20, 20), # decode

100 + (256, 0, 256), # prefill

101 + (1, 5, 5), # decode

102 + (27, 0, 27), # prefill

103 + (1, 4, 4), # decode

104 ],

105 expected_order=[0, 1, 6, 8, 4, 3, 2, 7, 5],

106 expected_modified=True,

107 ),

108 "new_request_single_token_prefill": ReorderTestCase(

109 requests=[

103 - (100, 0),

104 - (1, 0), # New request with only 1 token (STILL prefill)

105 - (50, 100),

106 - (1, 10),

110 + (100, 0, 100), # prefill

111 + (1, 0, 1), # prefill (single token, still prefill)

112 + (50, 100, 100), # long_extend

113 + (1, 10, 10), # decode

114 ],

108 - # Only index 3 is a true decode (has num_computed_tokens > 0)

115 expected_order=[3, 2, 0, 1],

116 expected_modified=True,

117 ),

118 "multiple_new_requests_single_token_prefill": ReorderTestCase(

119 requests=[

114 - (1, 0), # New prefill (1 token, no computed)

115 - (1, 0), # New prefill (1 token, no computed)

116 - (1, 50),

117 - (200, 0),

120 + (1, 0, 1), # prefill

121 + (1, 0, 1), # prefill

122 + (1, 50, 50), # decode

123 + (200, 0, 200), # prefill

124 ],

125 expected_order=[2, 1, 0, 3],

126 expected_modified=True,

127 ),

128 + "four_way_already_ordered": ReorderTestCase(

129 + requests=[

130 + (1, 100, 100), # decode

131 + (1, 50, 100), # short_extend

132 + (10, 50, 100), # long_extend

133 + (100, 0, 100), # prefill

134 + ],

135 + expected_order=[0, 1, 2, 3],

136 + expected_modified=False,

137 + ),

138 + "four_way_needs_reorder": ReorderTestCase(

139 + requests=[

140 + (100, 0, 100), # prefill

141 + (1, 50, 100), # short_extend

142 + (1, 100, 100), # decode

143 + (10, 50, 100), # long_extend

144 + ],

145 + expected_order=[2, 1, 3, 0],

146 + expected_modified=True,

147 + ),

148 + "four_way_multiple_short_extends": ReorderTestCase(

149 + requests=[

150 + (2, 100, 100), # decode

151 + (2, 50, 200), # short_extend

152 + (2, 75, 150), # short_extend

153 + (2, 200, 200), # decode

154 + ],

155 + expected_order=[0, 3, 2, 1],

156 + expected_modified=True,

157 + decode_threshold=2,

158 + ),

159 + "four_way_spec_decode_threshold": ReorderTestCase(

160 + requests=[

161 + (5, 100, 100), # decode

162 + (5, 50, 100), # short_extend

163 + (5, 0, 100), # prefill

164 + (10, 50, 100), # long_extend

165 + ],

166 + expected_order=[0, 1, 3, 2],

167 + expected_modified=True,

168 + decode_threshold=5,

169 + ),

170 }

171

172

@@ -129,8 +177,9 @@ def test_reorder_batch_to_split_decodes_and_prefills(test_case: ReorderTestCase)

177 req_ids = [f"r{i}" for i in range(len(test_case.requests))]

178 num_computed_tokens = np.array([r[1] for r in test_case.requests], dtype=np.int32)

179 num_scheduled_tokens = {f"r{i}": r[0] for i, r in enumerate(test_case.requests)}

180 + num_prompt_tokens = np.array([r[2] for r in test_case.requests], dtype=np.int32)

181

133 - input_batch = MockInputBatch(req_ids, num_computed_tokens)

182 + input_batch = MockInputBatch(req_ids, num_computed_tokens, num_prompt_tokens)

183 scheduler_output = MockSchedulerOutput(num_scheduled_tokens)

184

185 modified = reorder_batch_to_split_decodes_and_prefills(

tests/v1/e2e/test_hybrid_chunked_prefill.py

+2 -2

@@ -43,7 +43,7 @@ MESSAGES = [

43 pytest.param("Qwen/Qwen3.5-4B", marks=[large_gpu_mark(min_gb=40)]),

44 pytest.param(

45 "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8",

46 - marks=[large_gpu_mark(min_gb=80)] + multi_gpu_marks(num_gpus=2),

46 + marks=[large_gpu_mark(min_gb=80)] + multi_gpu_marks(num_gpus=4),

47 ),

48 ],

49 )

@@ -68,7 +68,7 @@ def test_mtp_speculative_mixed_batch_short_prefill(

68 max_num_batched_tokens=chunk_size,

69 max_model_len=512,

70 enforce_eager=True,

71 - tensor_parallel_size=2,

71 + tensor_parallel_size=4,

72 trust_remote_code=True,

73 enable_chunked_prefill=True,

74 enable_prefix_caching=enable_prefix_caching,

vllm/v1/attention/backend.py

@@ -362,6 +362,11 @@ class CommonAttentionMetadata:

362 dcp_local_seq_lens_cpu: torch.Tensor | None = None

363 """Sequence lengths of the local rank in decode context parallelism world"""

364

365 + is_prefilling: torch.Tensor | None = None

366 + """(batch_size,) bool tensor: True if request is still in prefill phase

367 + (num_computed_tokens < num_prompt_tokens). Used by some backends to

368 + distinguish actual decodes from short extends."""

369 +

370 # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)

371 _seq_lens_cpu: torch.Tensor | None = None

372 _num_computed_tokens_cpu: torch.Tensor | None = None

@@ -443,6 +448,7 @@ class CommonAttentionMetadata:

448 encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),

449 dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),

450 dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),

451 + is_prefilling=maybe_slice_reqs(self.is_prefilling),

452 )

453

454

vllm/v1/attention/backends/mamba_attn.py

+3 -1

@@ -358,7 +358,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

358

359 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

360 split_decodes_and_prefills(

361 - common_attn_metadata, decode_threshold=decode_threshold

361 + common_attn_metadata,

362 + decode_threshold=decode_threshold,

363 + treat_short_extends_as_decodes=False,

364 )

365 )

366

vllm/v1/attention/backends/utils.py

+49 -24

@@ -489,11 +489,15 @@ def split_decodes_and_prefills(

489 common_attn_metadata: CommonAttentionMetadata,

490 decode_threshold: int = 1,

491 require_uniform: bool = False,

492 + treat_short_extends_as_decodes: bool = True,

493 ) -> tuple[int, int, int, int]:

494 """

495 Assuming a reordered batch, finds the boundary between prefill and decode

496 requests.

497

498 + The batch is expected to be ordered as:

499 + decode → short_extend → long_extend → prefill

500 +

501 Args:

502 common_attn_metadata: CommonAttentionMetadata object containing the

503 batch metadata.

@@ -501,6 +505,9 @@ def split_decodes_and_prefills(

505 require_uniform: If True, requires that all decode requests have the

506 same query length. When set, some queries may be considered prefills

507 even if they are <= decode_threshold, in order to ensure uniformity.

508 + treat_short_extends_as_decodes: If True (default), short extends

509 + (query_len <= threshold but still prefilling) are counted as

510 + decodes. If False, they are counted as prefills.

511

512 Returns:

513 num_decodes: The number of decode requests.

@@ -513,8 +520,10 @@ def split_decodes_and_prefills(

520 num_tokens = common_attn_metadata.num_actual_tokens

521 query_start_loc = common_attn_metadata.query_start_loc_cpu

522

516 - if max_query_len <= decode_threshold and (

517 - not require_uniform or decode_threshold <= 1

523 + if (

524 + max_query_len <= decode_threshold

525 + and (not require_uniform or decode_threshold <= 1)

526 + and treat_short_extends_as_decodes

527 ):

528 return num_reqs, 0, num_tokens, 0

529

@@ -533,11 +542,14 @@ def split_decodes_and_prefills(

542 else:

543 is_prefill = query_lens > decode_threshold

544

545 + if not treat_short_extends_as_decodes:

546 + assert common_attn_metadata.is_prefilling is not None

547 + is_prefill |= common_attn_metadata.is_prefilling

548 +

549 if not torch.any(is_prefill):

550 return num_reqs, 0, num_tokens, 0

551

552 first_prefill = is_prefill.int().argmax(dim=-1).item()

540 - assert torch.all(query_lens[:first_prefill] <= decode_threshold)

553 num_decodes = first_prefill

554 num_prefills = num_reqs - num_decodes

555 num_decode_tokens = query_start_loc[first_prefill].item()

@@ -581,39 +593,52 @@ def reorder_batch_to_split_decodes_and_prefills(

593 Reorders the batch to split into prefill and decode requests; places all

594 requests with <= decode_threshold tokens at the front of the batch.

595

596 + The batch is reordered into 4 regions:

597 + decode: (num_scheduled <= threshold AND is not prefilling)

598 + short_extend: (num_scheduled <= threshold AND is chunked prefilling)

599 + long_extend: (num_scheduled > threshold AND is chunked prefilling)

600 + prefill: (num_computed == 0) # First chunks

601 +

602 Returns:

603 True if the batch was modified, False otherwise.

604 """

587 - # We now want to reorder the batch into decode → extend → prefill order

588 - # where:

589 - # decode: request with num_scheduled_tokens <= decode_threshold

590 - # extend: non-decode request with existing context

591 - # prefill: non-decode request with no existing context

592 - # NOTE for now we loosely use "decode" to mean requests where attention is

593 - # likely memory-bound and "prefill" to mean requests where attention is

594 - # likely compute-bound,

605 num_reqs = len(input_batch.req_ids)

606 num_scheduled_tokens = [

607 scheduler_output.num_scheduled_tokens[id] for id in input_batch.req_ids

608 ]

609 num_scheduled_tokens_np = np.array(num_scheduled_tokens)

610 num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]

601 -

602 - is_prefill = num_computed_tokens_np == 0

603 - is_decode = (num_scheduled_tokens_np <= decode_threshold) & (~is_prefill)

604 - is_extend = (num_scheduled_tokens_np > decode_threshold) & (~is_prefill)

605 -

606 - # Desired order: decode → extend → prefill

607 - req_regions = np.zeros(is_decode.shape, dtype=np.int32) # 0 = decode by default

608 - req_regions[is_extend] = 1

609 - req_regions[is_prefill] = 2

611 + num_prompt_tokens_np = input_batch.num_prompt_tokens[:num_reqs]

612 +

613 + has_context = num_computed_tokens_np > 0

614 + is_below_threshold = num_scheduled_tokens_np <= decode_threshold

615 + done_prefilling = num_computed_tokens_np >= num_prompt_tokens_np

616 +

617 + # Mutually exclusive categories (exactly one True per request):

618 + # 1. No context yet -> prefill

619 + # 2. Has context, above threshold -> long_extend

620 + # 3. Has context, below threshold, still prefilling -> short_extend

621 + # 4. Has context, below threshold, done prefilling -> decode

622 + is_pure_prefill = ~has_context

623 + is_long_extend = has_context & ~is_below_threshold

624 + is_short_extend = has_context & is_below_threshold & ~done_prefilling

625 + is_decode = has_context & is_below_threshold & done_prefilling

626 +

627 + # Desired order: decode → short_extend → long_extend → prefill

628 + req_regions = np.zeros(num_reqs, dtype=np.int32) # 0 = decode by default

629 + req_regions[is_short_extend] = 1

630 + req_regions[is_long_extend] = 2

631 + req_regions[is_pure_prefill] = 3

632

633 num_decodes = int(is_decode.sum())

612 - num_extends = int(is_extend.sum())

634 + num_short_extends = int(is_short_extend.sum())

635 + num_long_extends = int(is_long_extend.sum())

636 + num_prefills = int(is_pure_prefill.sum())

637

614 - target_regions = np.zeros(num_reqs, dtype=np.int32)

615 - target_regions[num_decodes : num_decodes + num_extends] = 1

616 - target_regions[num_decodes + num_extends :] = 2

638 + target_regions = np.repeat(

639 + [0, 1, 2, 3],

640 + [num_decodes, num_short_extends, num_long_extends, num_prefills],

641 + ).astype(np.int32)

642

643 needs_swap = req_regions != target_regions

644

vllm/v1/worker/gpu_input_batch.py

+7 -1

@@ -134,7 +134,13 @@ class InputBatch:

134 pin_memory=pin_memory,

135 )

136 self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy()

137 - self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)

137 + self.num_prompt_tokens_cpu_tensor = torch.zeros(

138 + (max_num_reqs,),

139 + device="cpu",

140 + dtype=torch.int32,

141 + pin_memory=pin_memory,

142 + )

143 + self.num_prompt_tokens = self.num_prompt_tokens_cpu_tensor.numpy()

144 self.num_computed_tokens_cpu_tensor = torch.zeros(

145 (max_num_reqs,),

146 device="cpu",

vllm/v1/worker/gpu_model_runner.py

+13 -28

@@ -740,19 +740,6 @@ class GPUModelRunner(

740

741 self.uniform_decode_query_len = 1 + self.num_spec_tokens

742

743 - # When spec decode is active, the mamba backend classifies requests

744 - # with query_len <= reorder_batch_threshold as "decodes". Prefill

745 - # chunks that fall under this threshold get processed via the decode

746 - # path, which stores intermediate states at sequential slots. We must

747 - # set num_accepted_tokens to the chunk's query_len for those requests

748 - # so the next iteration reads from the correct final-state slot.

749 - # Prefills that went through the actual prefill path should keep the

750 - # default value of 1 (the prefill path stores state at slot 0 only).

751 - self.needs_prefill_as_decode_slots: bool = False

752 - self.prefill_as_decode_num_tokens = self._make_buffer(

753 - self.max_num_reqs, dtype=torch.int32

754 - )

755 -

743 # Cudagraph dispatcher for runtime cudagraph dispatching.

744 self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

745

@@ -1369,16 +1356,6 @@ class GPUModelRunner(

1356 .int()

1357 .argmax(-1)

1358 )

1372 - spec_decode_active = bool(scheduler_output.scheduled_spec_decode_tokens)

1373 - if self.needs_prefill_as_decode_slots and spec_decode_active:

1374 - mamba_utils.update_accepted_tokens_for_prefill_as_decode(

1375 - self.input_batch,

1376 - self.prefill_as_decode_num_tokens,

1377 - self.num_accepted_tokens.gpu,

1378 - scheduler_output,

1379 - self.reorder_batch_threshold,

1380 - num_reqs,

1381 - )

1359

1360 if self.cache_config.mamba_cache_mode == "align":

1361 for i, num_tokens in enumerate(

@@ -1982,14 +1959,23 @@ class GPUModelRunner(

1959 attn_gid = self.routed_experts_attn_gid

1960 slot_mapping_attn = slot_mappings[attn_gid]

1961 self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()

1962 + # Compute is_prefilling: True if request is still in prefill phase

1963 + # (num_computed_tokens < num_prompt_tokens). Used by mamba backends to

1964 + # distinguish actual decodes from short extends.

1965 + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

1966 + :num_reqs_padded

1967 + ]

1968 + num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[

1969 + :num_reqs_padded

1970 + ]

1971 + is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu

1972 +

1973 cm_base = CommonAttentionMetadata(

1974 query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],

1975 query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],

1976 seq_lens=self.seq_lens.gpu[:num_reqs_padded],

1977 _seq_lens_cpu=self.seq_lens.cpu[:num_reqs_padded],

1990 - _num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[

1991 - :num_reqs_padded

1992 - ],

1978 + _num_computed_tokens_cpu=num_computed_tokens_cpu,

1979 num_reqs=num_reqs_padded,

1980 num_actual_tokens=num_tokens_padded,

1981 max_query_len=max_query_len,

@@ -1997,6 +1983,7 @@ class GPUModelRunner(

1983 block_table_tensor=block_table_gid_0,

1984 slot_mapping=slot_mapping_gid_0,

1985 causal=True,

1986 + is_prefilling=is_prefilling,

1987 )

1988

1989 if self.dcp_world_size > 1:

@@ -2048,8 +2035,6 @@ class GPUModelRunner(

2035 else 0

2036 )

2037

2051 - if isinstance(builder, Mamba2AttentionMetadataBuilder):

2052 - self.needs_prefill_as_decode_slots = True

2038 extra_attn_metadata_args = {}

2039 if use_spec_decode and isinstance(

2040 builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder)

vllm/v1/worker/mamba_utils.py

-42

@@ -266,45 +266,3 @@ def postprocess_mamba(

266 if src_block_idx == dest_block_idx:

267 num_accepted_tokens_cpu[i] = 1

268 do_mamba_copy_block(copy_bufs)

269 -

270 -

271 -def update_accepted_tokens_for_prefill_as_decode(

272 - input_batch: GPUInputBatch,

273 - prefill_as_decode_num_tokens: CpuGpuBuffer,

274 - num_accepted_tokens_gpu: torch.Tensor,

275 - scheduler_output: SchedulerOutput,

276 - decode_qlen_threshold: int | None,

277 - num_reqs: int,

278 -):

279 - """

280 - Adjusts num_accepted_tokens for prefill chunks processed via the decode path.

281 - This ensures subsequent iterations read from the correct sequential state slot

282 - instead of the default prefill slot 0. Not used by GDN attention, which manually

283 - separates short prefills and short decodes when building the attention metadata.

284 - """

285 - any_is_prefill = False

286 - for i in range(num_reqs):

287 - num_computed = input_batch.num_computed_tokens_cpu[i]

288 - num_prompt = input_batch.num_prompt_tokens[i]

289 - is_prefill = num_computed < num_prompt

290 - req_id = input_batch.req_ids[i]

291 - query_len = scheduler_output.num_scheduled_tokens[req_id]

292 -

293 - if is_prefill:

294 - classified_as_decode = (

295 - decode_qlen_threshold is not None and query_len <= decode_qlen_threshold

296 - )

297 - num_tokens = query_len if classified_as_decode else 1

298 - any_is_prefill = True

299 - else:

300 - num_tokens = -1

301 - prefill_as_decode_num_tokens.np[i] = num_tokens

302 -

303 - # We can skip the GPU transfer if there aren't any values to update

304 - if any_is_prefill:

305 - prefill_as_decode_num_tokens.copy_to_gpu(num_reqs)

306 - num_accepted_tokens_gpu[:num_reqs] = torch.where(

307 - prefill_as_decode_num_tokens.gpu[:num_reqs] != -1,

308 - prefill_as_decode_num_tokens.gpu[:num_reqs],

309 - num_accepted_tokens_gpu[:num_reqs],

310 - )