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

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 5, 17:15 UTC

[fix] Remove trtllm ragged mla prefills (#36540)

Signed-off-by: Olya Kozlova <okozlova@nvidia.com>

Olya Kozlova committed Mar 31, 2026 at 21:30 UTC 598190aac38a42d8c51ea46a3061e46d9078b3a5

8 files changed +185 -35

csrc/attention/merge_attn_states.cu

+48 -19

@@ -3,6 +3,7 @@

3 #include <ATen/cuda/CUDAContext.h>

4 #include <c10/cuda/CUDAGuard.h>

5 #include <algorithm>

6 +#include <limits>

8 #include "attention_dtypes.h"

9 #include "attention_utils.cuh"

@@ -17,7 +18,7 @@ __global__ void merge_attn_states_kernel(

18 const float* prefix_lse, const scalar_t* suffix_output,

19 const float* suffix_lse, const uint num_tokens, const uint num_heads,

20 const uint head_size, const uint prefix_head_stride,

20 - const uint output_head_stride) {

21 + const uint output_head_stride, const uint prefix_num_tokens) {

22 using pack_128b_t = uint4;

23 const uint pack_size = 16 / sizeof(scalar_t);

24 const uint threads_per_head = head_size / pack_size;

@@ -43,6 +44,22 @@ __global__ void merge_attn_states_kernel(

44 const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;

45 scalar_t* output_head_ptr = output + dst_head_offset;

47 + // If token_idx >= prefix_num_tokens, just copy from suffix

48 + if (token_idx >= prefix_num_tokens) {

49 + if (pack_offset < head_size) {

50 + pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(

51 + suffix_head_ptr)[pack_offset / pack_size];

52 + reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =

53 + s_out_pack;

54 + }

55 + if (output_lse != nullptr && pack_idx == 0) {

56 + float s_lse = suffix_lse[head_idx * num_tokens + token_idx];

57 + output_lse[head_idx * num_tokens + token_idx] = s_lse;

58 + }

59 + return;

60 + }

61 +

62 + // For tokens within prefix range, merge prefix and suffix

63 float p_lse = prefix_lse[head_idx * num_tokens + token_idx];

64 float s_lse = suffix_lse[head_idx * num_tokens + token_idx];

65 p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;

@@ -143,7 +160,8 @@ __global__ void merge_attn_states_kernel(

160 reinterpret_cast<float*>(prefix_lse.data_ptr()), \

161 reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \

162 reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens, \

146 - num_heads, head_size, prefix_head_stride, output_head_stride); \

163 + num_heads, head_size, prefix_head_stride, output_head_stride, \

164 + prefix_num_tokens); \

165 }

166

167 /*@brief Merges the attention states from prefix and suffix

@@ -157,14 +175,18 @@ __global__ void merge_attn_states_kernel(

175 * @param suffix_output [n,h,d] The suffix attention states.

176 * @param suffix_lse [h,n] The log-sum-exp values for the suffix attention

177 * states.

178 + * @param prefill_tokens_with_context Number of prefill tokens with context

179 + * For the first p tokens (0 <= token_idx < prefill_tokens_with_context), output

180 + * is computed by merging prefix_output and suffix_output. For remaining tokens

181 + * (prefill_tokens_with_context <= token_idx < n), output is copied directly

182 + * from suffix_output.

183 */

184 template <typename scalar_t>

162 -void merge_attn_states_launcher(torch::Tensor& output,

163 - std::optional<torch::Tensor> output_lse,

164 - const torch::Tensor& prefix_output,

165 - const torch::Tensor& prefix_lse,

166 - const torch::Tensor& suffix_output,

167 - const torch::Tensor& suffix_lse) {

185 +void merge_attn_states_launcher(

186 + torch::Tensor& output, std::optional<torch::Tensor> output_lse,

187 + const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,

188 + const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,

189 + const std::optional<int64_t> prefill_tokens_with_context) {

190 constexpr uint NUM_THREADS = 128;

191 const uint num_tokens = output.size(0);

192 const uint num_heads = output.size(1);

@@ -174,6 +196,14 @@ void merge_attn_states_launcher(torch::Tensor& output,

196 const uint pack_size = 16 / sizeof(scalar_t);

197 TORCH_CHECK(head_size % pack_size == 0,

198 "headsize must be multiple of pack_size:", pack_size);

199 +

200 + const uint prefix_num_tokens =

201 + prefill_tokens_with_context.has_value()

202 + ? static_cast<uint>(prefill_tokens_with_context.value())

203 + : num_tokens;

204 + TORCH_CHECK(prefix_num_tokens <= num_tokens,

205 + "prefix_num_tokens must be <= num_tokens");

206 +

207 float* output_lse_ptr = nullptr;

208 if (output_lse.has_value()) {

209 output_lse_ptr = output_lse.value().data_ptr<float>();

@@ -192,18 +222,17 @@ void merge_attn_states_launcher(torch::Tensor& output,

222 LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);

223 }

224

195 -#define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \

196 - { \

197 - merge_attn_states_launcher<scalar_t>(output, output_lse, prefix_output, \

198 - prefix_lse, suffix_output, \

199 - suffix_lse); \

225 +#define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \

226 + { \

227 + merge_attn_states_launcher<scalar_t>( \

228 + output, output_lse, prefix_output, prefix_lse, suffix_output, \

229 + suffix_lse, prefill_tokens_with_context); \

230 }

231

202 -void merge_attn_states(torch::Tensor& output,

203 - std::optional<torch::Tensor> output_lse,

204 - const torch::Tensor& prefix_output,

205 - const torch::Tensor& prefix_lse,

206 - const torch::Tensor& suffix_output,

207 - const torch::Tensor& suffix_lse) {

232 +void merge_attn_states(

233 + torch::Tensor& output, std::optional<torch::Tensor> output_lse,

234 + const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,

235 + const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,

236 + std::optional<int64_t> prefill_tokens_with_context = std::nullopt) {

237 DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);

238 }

csrc/ops.h

+5 -6

@@ -53,12 +53,11 @@ void paged_attention_v2(

53 const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,

54 const int64_t blocksparse_head_sliding_step);

56 -void merge_attn_states(torch::Tensor& output,

57 - std::optional<torch::Tensor> output_lse,

58 - const torch::Tensor& prefix_output,

59 - const torch::Tensor& prefix_lse,

60 - const torch::Tensor& suffix_output,

61 - const torch::Tensor& suffix_lse);

56 +void merge_attn_states(

57 + torch::Tensor& output, std::optional<torch::Tensor> output_lse,

58 + const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,

59 + const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,

60 + const std::optional<int64_t> prefill_tokens_with_context);

61 #ifndef USE_ROCM

62 void convert_vertical_slash_indexes(

63 torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]

csrc/torch_bindings.cpp

+2 -1

@@ -73,7 +73,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

73 " Tensor prefix_output,"

74 " Tensor prefix_lse,"

75 " Tensor suffix_output,"

76 - " Tensor suffix_lse) -> ()");

76 + " Tensor suffix_lse,"

77 + " int!? prefill_tokens_with_context) -> ()");

78 ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);

79 #ifndef USE_ROCM

80 ops.def(

tests/kernels/attention/test_merge_attn_states.py

+26 -2

@@ -20,7 +20,11 @@ def merge_attn_states_torch(

20 suffix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]

21 suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]

22 output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS]

23 + prefill_tokens_with_context: int | None = None,

24 ):

25 + # Apply prefill_tokens_with_context mask if needed

26 + if prefill_tokens_with_context is None:

27 + prefill_tokens_with_context = output.shape[0]

28 p_lse = prefix_lse

29 s_lse = suffix_lse

30 # inf -> -inf

@@ -28,6 +32,9 @@ def merge_attn_states_torch(

32 s_lse[s_lse == torch.inf] = -torch.inf

33 # max_lse [NUM_HEADS, NUM_TOKENS]

34 max_lse = torch.maximum(p_lse, s_lse)

35 +

36 + mask = torch.ones((prefix_lse.shape[1], 1, 1), device=p_lse.device)

37 + mask[prefill_tokens_with_context:].fill_(0)

38 p_lse = p_lse - max_lse

39 s_lse = s_lse - max_lse

40 p_lse_exp = torch.exp(p_lse)

@@ -35,11 +42,16 @@ def merge_attn_states_torch(

42 out_se = p_lse_exp + s_lse_exp

43 if output_lse is not None:

44 output_lse = torch.log(out_se) + max_lse

45 + output_lse[prefill_tokens_with_context:] = suffix_lse[

46 + prefill_tokens_with_context:

47 + ]

48 p_scale = p_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]

49 s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]

50 p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]

51 s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]

42 - output = prefix_output * p_scale + suffix_output * s_scale

52 + output.copy_(

53 + prefix_output * p_scale * mask + suffix_output * (s_scale * mask + (1 - mask))

54 + )

55 return output, output_lse

@@ -90,13 +102,18 @@ def generate_markdown_table():

102 )

103

104

105 +@pytest.mark.parametrize("prefill_tokens_with_context", [None, 128])

106 @pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)

107 @pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)

108 @pytest.mark.parametrize("head_size", HEAD_SIZES)

109 @pytest.mark.parametrize("output_dtype", DTYPES)

110 @torch.inference_mode()

111 def test_merge_attn_states(

99 - num_tokens: int, num_query_heads: int, head_size: int, output_dtype: torch.dtype

112 + prefill_tokens_with_context: int | None,

113 + num_tokens: int,

114 + num_query_heads: int,

115 + head_size: int,

116 + output_dtype: torch.dtype,

117 ):

118 if not current_platform.is_cuda():

119 pytest.skip(

@@ -111,6 +128,7 @@ def test_merge_attn_states(

128 print(

129 f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "

130 f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, "

131 + f"prefill_tokens_with_context: {prefill_tokens_with_context}, "

132 f"Device: {current_platform.get_device_name()}"

133 )

134

@@ -164,6 +182,7 @@ def test_merge_attn_states(

182 suffix_output,

183 suffix_lse_torch,

184 output_lse_torch,

185 + prefill_tokens_with_context,

186 )

187 torch.accelerator.synchronize()

188

@@ -176,6 +195,7 @@ def test_merge_attn_states(

195 suffix_output,

196 suffix_lse_torch,

197 output_lse_torch,

198 + prefill_tokens_with_context,

199 )

200 end.record()

201 torch.accelerator.synchronize()

@@ -199,6 +219,7 @@ def test_merge_attn_states(

219 suffix_output,

220 suffix_lse,

221 output_lse_ref_triton,

222 + prefill_tokens_with_context,

223 )

224 torch.accelerator.synchronize()

225

@@ -211,6 +232,7 @@ def test_merge_attn_states(

232 suffix_output,

233 suffix_lse,

234 output_lse_ref_triton,

235 + prefill_tokens_with_context,

236 )

237 end.record()

238 torch.accelerator.synchronize()

@@ -231,6 +253,7 @@ def test_merge_attn_states(

253 suffix_output,

254 suffix_lse,

255 output_lse_cuda,

256 + prefill_tokens_with_context,

257 )

258 torch.accelerator.synchronize()

259

@@ -243,6 +266,7 @@ def test_merge_attn_states(

266 suffix_output,

267 suffix_lse,

268 output_lse_cuda,

269 + prefill_tokens_with_context,

270 )

271 end.record()

272 torch.accelerator.synchronize()

vllm/_custom_ops.py

+8 -1

@@ -264,9 +264,16 @@ def merge_attn_states(

264 suffix_output: torch.Tensor,

265 suffix_lse: torch.Tensor,

266 output_lse: torch.Tensor | None = None,

267 + prefill_tokens_with_context: int | None = None,

268 ) -> None:

269 torch.ops._C.merge_attn_states(

269 - output, output_lse, prefix_output, prefix_lse, suffix_output, suffix_lse

270 + output,

271 + output_lse,

272 + prefix_output,

273 + prefix_lse,

274 + suffix_output,

275 + suffix_lse,

276 + prefill_tokens_with_context,

277 )

278

279

vllm/model_executor/layers/attention/mla_attention.py

+14 -2

@@ -1181,6 +1181,7 @@ class MLACommonPrefillMetadata:

1181 padded_local_cu_seq_lens: torch.Tensor | None = None

1182 cu_seq_lens_lst: list[list[int]] | None = None

1183 chunk_size: int | None = None

1184 + prefill_tokens_with_context: int | None = None

1185

1186 block_table: torch.Tensor

1187 query_start_loc: torch.Tensor

@@ -1743,6 +1744,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1744 prefill_query_start_loc = (

1745 query_start_loc[reqs_start:] - query_start_loc[reqs_start]

1746 )

1747 + prefill_query_start_loc_cpu = (

1748 + query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]

1749 + )

1750

1751 chunked_context_metadata = None

1752 if max_context_len_cpu > 0:

@@ -1864,6 +1868,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1868 if self._use_cudnn_prefill

1869 else MLACommonPrefillMetadata.ChunkedContextMetadata

1870 )

1871 + prefill_tokens_with_context = None

1872 + if num_prefills_with_context_cpu > 0:

1873 + prefill_tokens_with_context = prefill_query_start_loc_cpu[

1874 + num_prefills_with_context_cpu

1875 + ].item()

1876 if self.dcp_world_size > 1:

1877 chunked_context_metadata = chunked_context_metadata_cls(

1878 cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

@@ -1883,6 +1892,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1892 ),

1893 cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),

1894 chunk_size=padded_local_max_context_chunk_across_ranks,

1895 + prefill_tokens_with_context=prefill_tokens_with_context,

1896 )

1897 else:

1898 chunked_context_metadata = chunked_context_metadata_cls(

@@ -1896,6 +1906,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1906 ),

1907 chunk_total_token=chunk_total_token,

1908 workspace=self.chunked_prefill_workspace,

1909 + prefill_tokens_with_context=prefill_tokens_with_context,

1910 )

1911

1912 if self._use_cudnn_prefill:

@@ -2382,14 +2393,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2393 assert prefill.chunked_context.seq_lens[chunk_idx] is not None

2394 assert prefill.workspace_buffer is not None

2395

2385 - out = torch.zeros(

2396 + out = torch.empty(

2397 q.shape[0],

2398 q.shape[1],

2399 v.shape[2],

2400 device=q.device,

2401 dtype=prefill.output_dtype,

2402 )

2392 - prefill.workspace_buffer.fill_(0)

2403

2404 attn_out, lse = trtllm_ragged_attention_deepseek(

2405 query=q,

@@ -2691,6 +2701,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2701 )

2702

2703 if has_context:

2704 + assert prefill_metadata.chunked_context is not None

2705 suffix_output, suffix_lse = output_prefill

2706 if self.dcp_world_size > 1:

2707 context_output, context_lse = (

@@ -2719,6 +2730,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2730 prefix_lse=context_lse,

2731 suffix_output=suffix_output,

2732 suffix_lse=suffix_lse,

2733 + prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context,

2734 )

2735 else:

2736 output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2)

vllm/v1/attention/ops/merge_attn_states.py

+43 -2

@@ -13,7 +13,36 @@ def merge_attn_states(

13 suffix_output: torch.Tensor,

14 suffix_lse: torch.Tensor,

15 output_lse: torch.Tensor | None = None,

16 + prefill_tokens_with_context: int | None = None,

17 ) -> None:

18 + """Merge partial attention outputs from prefix (KV cache) and suffix

19 + (new tokens) into a single output tensor using the log-sum-exp (LSE)

20 + rescaling method described in section 2.2 of

21 + https://www.arxiv.org/pdf/2501.01005.

22 +

23 + For tokens that have prefix context (token index < prefill_tokens_with_context),

24 + the prefix and suffix partial outputs are combined as a weighted sum.

25 + For tokens without prefix context, the suffix output is copied directly.

26 +

27 + Args:

28 + output: Output tensor of shape [NUM_TOKENS, NUM_HEADS, HEAD_SIZE].

29 + prefix_output: Partial attention output over the prefix (KV cache),

30 + shape [NUM_TOKENS, NUM_HEADS, HEAD_SIZE].

31 + prefix_lse: Log-sum-exp values for the prefix attention,

32 + shape [NUM_HEADS, NUM_TOKENS].

33 + suffix_output: Partial attention output over the suffix (new KV),

34 + shape [NUM_TOKENS, NUM_HEADS, HEAD_SIZE].

35 + suffix_lse: Log-sum-exp values for the suffix attention,

36 + shape [NUM_HEADS, NUM_TOKENS].

37 + output_lse: Optional tensor to store the merged LSE values,

38 + shape [NUM_HEADS, NUM_TOKENS]. If None, LSE is not written out.

39 + prefill_tokens_with_context: Number of prefill tokens that have

40 + prefix context and therefore require merging. Tokens at indices

41 + >= this value are decode or context-free prefill tokens whose

42 + output is taken directly from suffix_output. If None, all tokens

43 + are treated as having context.

44 + """

45 +

46 # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel

47 # does not support FP8 dtype, fallback to use Triton kernel.

48 def supported_dtypes(o: torch.Tensor) -> bool:

@@ -37,11 +66,23 @@ def merge_attn_states(

66 from vllm._custom_ops import merge_attn_states

68 return merge_attn_states(

40 - output, prefix_output, prefix_lse, suffix_output, suffix_lse, output_lse

69 + output,

70 + prefix_output,

71 + prefix_lse,

72 + suffix_output,

73 + suffix_lse,

74 + output_lse,

75 + prefill_tokens_with_context,

76 )

77 else:

78 from vllm.v1.attention.ops.triton_merge_attn_states import merge_attn_states

80 return merge_attn_states(

46 - output, prefix_output, prefix_lse, suffix_output, suffix_lse, output_lse

81 + output,

82 + prefix_output,

83 + prefix_lse,

84 + suffix_output,

85 + suffix_lse,

86 + output_lse,

87 + prefill_tokens_with_context,

88 )

vllm/v1/attention/ops/triton_merge_attn_states.py

+39 -2

@@ -15,6 +15,7 @@ def merge_attn_states(

15 suffix_output: torch.Tensor,

16 suffix_lse: torch.Tensor,

17 output_lse: torch.Tensor | None = None,

18 + prefill_tokens_with_context: int | None = None,

19 ) -> None:

20 num_tokens = output.shape[0]

21 num_query_heads = output.shape[1]

@@ -25,6 +26,11 @@ def merge_attn_states(

26 # backend.

27 prefix_head_stride = prefix_output.stride(1)

28 output_head_stride = output.stride(1)

29 +

30 + # If prefill_tokens_with_context is None, all tokens should use prefix context

31 + if prefill_tokens_with_context is None:

32 + prefill_tokens_with_context = num_tokens

33 +

34 # TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead.

35 merge_attn_states_kernel[(num_tokens, num_query_heads)](

36 output,

@@ -38,6 +44,7 @@ def merge_attn_states(

44 head_size,

45 padded_head_size,

46 output_lse is not None,

47 + prefill_tokens_with_context,

48 )

@@ -54,12 +61,44 @@ def merge_attn_states_kernel(

61 HEAD_SIZE: tl.constexpr,

62 PADDED_HEAD_SIZE: tl.constexpr,

63 OUTPUT_LSE: tl.constexpr,

64 + prefill_tokens_with_context: tl.constexpr,

65 ):

66 token_idx = tl.program_id(0)

67 num_tokens = tl.num_programs(0)

68 head_idx = tl.program_id(1)

69 num_heads = tl.num_programs(1)

71 + prefix_mask = token_idx < prefill_tokens_with_context

72 +

73 + head_arange = tl.arange(0, PADDED_HEAD_SIZE)

74 + head_mask = head_arange < HEAD_SIZE

75 +

76 + # For tokens without context (token_idx >= prefill_tokens_with_context),

77 + # directly copy from suffix_output

78 + if not prefix_mask:

79 + s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)

80 + if OUTPUT_LSE:

81 + tl.store(output_lse + head_idx * num_tokens + token_idx, s_lse)

82 +

83 + s_out = tl.load(

84 + suffix_output

85 + + token_idx * num_heads * prefix_head_stride

86 + + head_idx * prefix_head_stride

87 + + head_arange,

88 + mask=head_mask,

89 + )

90 + tl.store(

91 + output

92 + + token_idx * num_heads * output_head_stride

93 + + head_idx * output_head_stride

94 + + head_arange,

95 + s_out,

96 + mask=head_mask,

97 + )

98 + return

99 +

100 + # For tokens with context (token_idx < prefill_tokens_with_context),

101 + # perform normal merge operation

102 p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)

103 s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)

104

@@ -83,8 +122,6 @@ def merge_attn_states_kernel(

122 out_lse = tl.log(out_se) + max_lse

123 tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)

124

86 - head_arange = tl.arange(0, PADDED_HEAD_SIZE)

87 - head_mask = head_arange < HEAD_SIZE

125 p_out = tl.load(

126 prefix_output

127 + token_idx * num_heads * prefix_head_stride