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

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 8, 19:00 UTC

[Perf] Disable clean_logits in deepgemm fp8_mqa_logits kernel (#33568)

Xin Yang committed Feb 5, 2026 at 17:34 UTC 79028d438859162841d35bbf2a91eaa8236733fa

4 files changed +61 -27

tests/kernels/attention/test_deepgemm_attention.py

+13 -7

@@ -95,7 +95,8 @@ def _ref_fp8_mqa_logits(

95 @pytest.mark.skipif(

96 not current_platform.has_device_capability(90), reason="SM90 and SM100 only"

97 )

98 -def test_deepgemm_fp8_mqa_logits():

98 +@pytest.mark.parametrize("clean_logits", [True, False])

99 +def test_deepgemm_fp8_mqa_logits(clean_logits: bool):

100 torch.manual_seed(0)

101 random.seed(0)

102 num_heads, head_dim = 32, 128

@@ -126,7 +127,9 @@ def test_deepgemm_fp8_mqa_logits():

127

128 q_fp8 = q.to(torch.float8_e4m3fn)

129 kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)

129 - logits = fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)

130 + logits = fp8_mqa_logits(

131 + q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits

132 + )

133

134 ref_logits = _ref_fp8_mqa_logits(

135 q=q,

@@ -135,13 +138,14 @@ def test_deepgemm_fp8_mqa_logits():

138 cu_seqlen_ks=ks,

139 cu_seqlen_ke=ke,

140 )

138 -

141 ref_neginf_mask = ref_logits == float("-inf")

140 - neginf_mask = logits == float("-inf")

141 - assert torch.equal(neginf_mask, ref_neginf_mask)

142 +

143 + if clean_logits:

144 + neginf_mask = logits == float("-inf")

145 + assert torch.equal(neginf_mask, ref_neginf_mask)

146

147 ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)

144 - logits = logits.masked_fill(neginf_mask, 0)

148 + logits = logits.masked_fill(ref_neginf_mask, 0)

149 diff = calc_diff(logits, ref_logits)

150 assert diff < 1e-3, f"{diff=}"

151

@@ -201,7 +205,8 @@ def _ref_fp8_paged_mqa_logits(

205 @pytest.mark.skipif(

206 not current_platform.has_device_capability(90), reason="SM90 and SM100 only"

207 )

204 -def test_deepgemm_fp8_paged_mqa_logits():

208 +@pytest.mark.parametrize("clean_logits", [True, False])

209 +def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):

210 torch.manual_seed(0)

211 random.seed(0)

212

@@ -264,6 +269,7 @@ def test_deepgemm_fp8_paged_mqa_logits():

269 block_tables,

270 schedule_metadata,

271 max_model_len,

272 + clean_logits=clean_logits,

273 )

274

275 ref_logits = _ref_fp8_paged_mqa_logits(

tests/kernels/test_top_k_per_row.py

+38 -18

@@ -6,6 +6,7 @@ import pytest

6 import torch

8 from vllm.platforms import current_platform

9 +from vllm.utils.torch_utils import set_random_seed

11 # Test parameters

12 NUM_ROWS = [1, 32, 2050]

@@ -20,6 +21,7 @@ def create_random_logits(

21 row_ends: torch.Tensor,

22 dtype: torch.dtype,

23 seed: int,

24 + clean_logits: bool,

25 data_generation: str,

26 ) -> torch.Tensor:

27 """Create random logits tensor for testing."""

@@ -48,8 +50,9 @@ def create_random_logits(

50 )

51 logits = logits_bits.view(dtype)

51 - for i, end in enumerate(row_ends):

52 - logits[i, end:] = float("-inf")

53 + if clean_logits:

54 + for i, end in enumerate(row_ends):

55 + logits[i, end:] = float("-inf")

56 return logits

@@ -121,21 +124,26 @@ def compare_top_k_results(

124

125 @pytest.mark.parametrize("num_rows", NUM_ROWS)

126 @pytest.mark.parametrize("top_k", TOP_K_VALUES)

127 +@pytest.mark.parametrize("clean_logits", [True, False])

128 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

129 @torch.inference_mode()

130 def test_top_k_per_row(

131 num_rows: int,

132 top_k: int,

133 + clean_logits: bool,

134 ) -> None:

135 """

136 Test top_k_per_row.

137 """

138 + set_random_seed(0)

139 torch.set_default_device("cuda:0")

140

141 # Create test data

142 vocab_size = 20000

143 row_starts, row_ends = create_row_boundaries(num_rows, vocab_size)

138 - logits = create_random_logits(row_starts, row_ends, torch.float32, 42, "random")

144 + logits = create_random_logits(

145 + row_starts, row_ends, torch.float32, 42, clean_logits, "random"

146 + )

147

148 # Create output tensors

149 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")

@@ -153,11 +161,12 @@ def test_top_k_per_row(

161 )

162

163 # Run reference implementation

156 - torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)[1]

157 - mask_lo = torch_indices >= 0

158 - mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0

159 - mask = mask_lo & mask_hi

160 - torch_indices = torch_indices.masked_fill(~mask, -1)

164 + torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")

165 + for i in range(num_rows):

166 + row_end = int(row_ends[i])

167 + k_i = min(top_k, row_end)

168 + idx = logits[i, :row_end].topk(k_i, dim=-1)[1]

169 + torch_indices[i, :k_i] = idx

170

171 # Compare results

172 assert compare_top_k_results(

@@ -170,6 +179,7 @@ def _run_top_k_per_row_decode_test(

179 batch_size: int,

180 next_n: int,

181 vocab_size: int,

182 + clean_logits: bool,

183 data_generation: str,

184 ) -> None:

185 """

@@ -180,14 +190,18 @@ def _run_top_k_per_row_decode_test(

190 # Create test data

191 num_rows = batch_size * next_n

192 seq_lens = torch.randint(

183 - vocab_size, (batch_size,), dtype=torch.int32, device="cuda"

193 + low=next_n,

194 + high=vocab_size,

195 + size=(batch_size,),

196 + dtype=torch.int32,

197 + device="cuda",

198 )

199 row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")

200 row_indices = torch.arange(num_rows, device="cuda") // next_n

201 next_n_offset = torch.arange(num_rows, device="cuda") % next_n

202 row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1

203 logits = create_random_logits(

190 - row_starts, row_ends, torch.float32, 42, data_generation

204 + row_starts, row_ends, torch.float32, 42, clean_logits, data_generation

205 )

206

207 # Create output tensors

@@ -208,11 +222,12 @@ def _run_top_k_per_row_decode_test(

222 torch.cuda.synchronize()

223

224 # Run reference implementation

211 - torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)[1]

212 - mask_lo = torch_indices >= 0

213 - mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0

214 - mask = mask_lo & mask_hi

215 - torch_indices = torch_indices.masked_fill(~mask, -1)

225 + torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")

226 + for i in range(num_rows):

227 + row_end = int(row_ends[i])

228 + k_i = min(top_k, row_end)

229 + idx = logits[i, :row_end].topk(k_i, dim=-1)[1]

230 + torch_indices[i, :k_i] = idx

231

232 # Compare results

233 assert compare_top_k_results(

@@ -223,6 +238,7 @@ def _run_top_k_per_row_decode_test(

238 @pytest.mark.parametrize("top_k", TOP_K_VALUES)

239 @pytest.mark.parametrize("batch_size", BATCH_SIZE)

240 @pytest.mark.parametrize("next_n", NEXT_N)

241 +@pytest.mark.parametrize("clean_logits", [True, False])

242 @pytest.mark.parametrize("data_generation", DATA_GENERATION)

243 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

244 @torch.inference_mode()

@@ -230,28 +246,32 @@ def test_top_k_per_row_decode(

246 top_k: int,

247 batch_size: int,

248 next_n: int,

249 + clean_logits: bool,

250 data_generation: str,

251 ) -> None:

252 """

253 Test top_k_per_row with seq_lens tensor.

254 """

255 + set_random_seed(0)

256 vocab_size = 20000

257 _run_top_k_per_row_decode_test(

240 - top_k, batch_size, next_n, vocab_size, data_generation

258 + top_k, batch_size, next_n, vocab_size, clean_logits, data_generation

259 )

260

261

262 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

263 +@pytest.mark.parametrize("clean_logits", [True, False])

264 @torch.inference_mode()

246 -def test_top_k_per_row_decode_large_vocab_size() -> None:

265 +def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:

266 """

267 Test top_k_per_row_decode with large vocabulary size.

268 """

269 + set_random_seed(0)

270 top_k = 2048

271 batch_size = 2

272 next_n = 2

273 vocab_size = 300000

274 data_generation = "random"

275 _run_top_k_per_row_decode_test(

256 - top_k, batch_size, next_n, vocab_size, data_generation

276 + top_k, batch_size, next_n, vocab_size, clean_logits, data_generation

277 )

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -108,6 +108,7 @@ def sparse_attn_indexer(

108 weights[chunk.token_start : chunk.token_end],

109 chunk.cu_seqlen_ks,

110 chunk.cu_seqlen_ke,

111 + clean_logits=False,

112 )

113 num_rows = logits.shape[0]

114

@@ -157,6 +158,7 @@ def sparse_attn_indexer(

158 decode_metadata.block_table,

159 decode_metadata.schedule_metadata,

160 max_model_len=max_model_len,

161 + clean_logits=False,

162 )

163

164 num_rows = logits.shape[0]

vllm/utils/deep_gemm.py

+8 -2

@@ -242,6 +242,7 @@ def fp8_mqa_logits(

242 weights: torch.Tensor,

243 cu_seqlen_ks: torch.Tensor,

244 cu_seqlen_ke: torch.Tensor,

245 + clean_logits: bool,

246 ) -> torch.Tensor:

247 """Compute FP8 MQA logits for a single sequence without KV paging.

248

@@ -256,6 +257,7 @@ def fp8_mqa_logits(

257 shape [M], dtype int32.

258 cu_seqlen_ke: End indices (exclusive) for valid K per query position,

259 shape [M], dtype int32.

260 + clean_logits: Whether to clean the unfilled logits into `-inf`.

261

262 Returns:

263 Logits tensor of shape [M, N], dtype `torch.float32`.

@@ -263,7 +265,9 @@ def fp8_mqa_logits(

265 _lazy_init()

266 if _fp8_mqa_logits_impl is None:

267 return _missing()

266 - return _fp8_mqa_logits_impl(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)

268 + return _fp8_mqa_logits_impl(

269 + q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=clean_logits

270 + )

271

272

273 def get_paged_mqa_logits_metadata(

@@ -295,6 +299,7 @@ def fp8_paged_mqa_logits(

299 block_tables: torch.Tensor,

300 schedule_metadata: torch.Tensor,

301 max_model_len: int,

302 + clean_logits: bool,

303 ) -> torch.Tensor:

304 """Compute FP8 MQA logits using paged KV-cache.

305

@@ -312,6 +317,7 @@ def fp8_paged_mqa_logits(

317 schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;

318 used to distribute work across SMs.

319 max_model_len: Maximum sequence length used to size the logits output.

320 + clean_logits: Whether to clean the unfilled logits into `-inf`.

321

322 Returns:

323 Logits tensor of shape [B * next_n, max_model_len], dtype

@@ -328,7 +334,7 @@ def fp8_paged_mqa_logits(

334 block_tables,

335 schedule_metadata,

336 max_model_len,

331 - clean_logits=True,

337 + clean_logits=clean_logits,

338 )

339

340