@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fd4a151262

[Kernel] Reuse Qwen4Exp HC combine-norm for MTP input (#54687)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

Canlin Guo committed Sep 4, 2026 at 17:57 UTC fd4a1512628ad17944095263c1fe89598710a3ce

5 files changed +93 -40

tests/models/qwen4_exp/test_hc_ops.py

+33

@@ -68,6 +68,18 @@ def test_hc_combine() -> None:

68 torch.testing.assert_close(actual, expected.flatten(-2).to(torch.bfloat16))

71 +def test_hc_combine_unit_injection() -> None:

72 + torch.manual_seed(0)

73 + block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")

74 + residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")

75 +

76 + actual = hc_combine(residual, block_output, None, HC)

77 + expected = residual.unflatten(-1, (HC, HIDDEN_SIZE))

78 + expected = expected + block_output.unsqueeze(-2)

79 +

80 + assert torch.equal(actual, expected.flatten(-2))

81 +

82 +

83 def test_hc_combine_norm() -> None:

84 torch.manual_seed(0)

85 block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")

@@ -92,3 +104,24 @@ def test_hc_combine_norm() -> None:

104

105 torch.testing.assert_close(actual, expected)

106 torch.testing.assert_close(actual_norm, expected_norm.to(torch.bfloat16))

107 +

108 +

109 +@pytest.mark.parametrize("num_tokens", [1, 17, 2048])

110 +def test_hc_combine_norm_unit_injection(num_tokens: int) -> None:

111 + torch.manual_seed(0)

112 + embedding = torch.randn(

113 + num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda"

114 + )

115 + hidden = torch.randn(

116 + num_tokens, HC, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda"

117 + )

118 + weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")

119 +

120 + actual, actual_norm = hc_combine_norm(

121 + hidden.flatten(1), embedding, None, weight, EPS, HC

122 + )

123 +

124 + expected = (hidden + embedding.unsqueeze(1)).flatten(1)

125 + expected_norm = grouped_gemma_rmsnorm(expected, weight, EPS, HC)

126 + assert torch.equal(actual, expected)

127 + torch.testing.assert_close(actual_norm, expected_norm)

vllm/models/qwen4_exp/nvidia/hyperconnection.py

+8 -6

@@ -52,9 +52,10 @@ class GatedResidual(nn.Module):

53 ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged

54 low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean

55 - over the HC streams). When passed a pending block output and an injection,

56 - it fuses their residual combine with the RMSNorm. Final mixers use

57 - ``use_combine=False`` and do not produce a new injection.

55 + over the HC streams). When passed a pending block output, it fuses its

56 + residual combine with the RMSNorm. A missing injection selects unit-weight

57 + combine. Final mixers use ``use_combine=False`` and do not produce a new

58 + injection.

60 Weights: the norm owns the grouped GemmaRMSNorm affine; the projections

61 are vLLM Linear modules (merged replicated linear for down+inject), so

@@ -153,13 +154,14 @@ class GatedResidual(nn.Module):

154 self,

155 hidden_states: torch.Tensor,

156 prev_block_output: torch.Tensor,

156 - prev_injection: torch.Tensor,

157 + prev_injection: torch.Tensor | None,

158 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

159 """Consume a pending combine, then prepare the next block input.

160

161 ``hidden_states`` is the multi-stream state from before the pending

162 block's mix. Its combine with ``block_output`` is fused with this

162 - module's input RMSNorm.

163 + module's input RMSNorm. A missing injection applies the block output

164 + to every stream with unit weight.

165 """

166 hidden_states, xn = hc_combine_norm(

167 hidden_states,

@@ -189,7 +191,7 @@ class GatedResidual(nn.Module):

191 self,

192 hidden_states: torch.Tensor,

193 block_output: torch.Tensor,

192 - injection: torch.Tensor,

194 + injection: torch.Tensor | None,

195 ) -> torch.Tensor:

196 return hc_combine(hidden_states, block_output, injection, self.hc_count)

197

vllm/models/qwen4_exp/nvidia/model.py

+4 -2

@@ -284,11 +284,13 @@ class Qwen4ExpDecoderLayer(nn.Module):

284 query_start_loc: torch.Tensor | None,

285 ngram_context: torch.Tensor | None,

286 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

287 + if prev_block_output is None:

288 + assert prev_injection is None

289 attn_hc = self.attn_hyper_connection

290 if self.ple is not None:

291 # PLE adds directly to the multi-stream state, so pending HC state

292 # must be materialized before the addition.

291 - if prev_block_output is not None and prev_injection is not None:

293 + if prev_block_output is not None:

294 hidden_states = attn_hc.combine(

295 hidden_states, prev_block_output, prev_injection

296 )

@@ -304,7 +306,7 @@ class Qwen4ExpDecoderLayer(nn.Module):

306 )

307

308 # Fuse a pending combine with this HC module's mix when possible.

307 - if prev_block_output is not None and prev_injection is not None:

309 + if prev_block_output is not None:

310 hidden_states, block_input, injection = attn_hc.combine_and_mix(

311 hidden_states, prev_block_output, prev_injection

312 )

vllm/models/qwen4_exp/nvidia/mtp.py

+3 -4

@@ -280,6 +280,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):

280 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:

281 hc_count = self.hc_count

282 hidden_size = self.hidden_size

283 + prev_block_output: torch.Tensor | None = None

284

285 if get_pp_group().is_first_rank:

286 assert hidden_states is not None

@@ -300,10 +301,8 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):

301 num_tokens, hc_count, hidden_size

302 )

303 hidden_states = self.fc_hidden(hidden_states)

303 - # Add the embedding residual to every branch, then fold back

304 - # to [T, hc_count*H] (HC outer, HS inner) for the HC decoder.

305 - hidden_states = inputs_embeds.unsqueeze(-2) + hidden_states

304 hidden_states = hidden_states.flatten(-2)

305 + prev_block_output = inputs_embeds

306 else:

307 assert intermediate_tensors is not None

308 hidden_states = intermediate_tensors["hidden_states"]

@@ -312,7 +311,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):

311 layer = self.layers[current_step_idx]

312 hidden_states, block_output, injection = layer(

313 hidden_states=hidden_states,

315 - prev_block_output=None,

314 + prev_block_output=prev_block_output,

315 prev_injection=None,

316 positions=positions,

317 input_ids=None,

vllm/models/qwen4_exp/nvidia/ops/hc.py

+45 -28

@@ -215,14 +215,17 @@ def _hc_combine_kernel(

215 if launch_pdl:

216 tl.extra.cuda.gdc_wait()

217

218 - inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)

218 + if inj_ptr is not None:

219 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)

220 block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)

221 res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0)

222

223 # Keeping HC as a broadcast dimension is faster here than four separate

224 # residual load/store sequences.

224 - inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)

225 - out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None]

225 + if inj_ptr is not None:

226 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)

227 + block = block.to(tl.float32)[None, :] * inj[:, None]

228 + out = res.to(tl.float32) + block.to(tl.float32)

229

230 if launch_pdl:

231 tl.extra.cuda.gdc_launch_dependents()

@@ -232,17 +235,20 @@ def _hc_combine_kernel(

235 def _hc_combine(

236 residual: torch.Tensor,

237 block_output: torch.Tensor,

235 - injection_logits: torch.Tensor,

238 + injection_logits: torch.Tensor | None,

239 hc_count: int,

240 ) -> torch.Tensor:

241 N, DIM = residual.shape

242 assert DIM % hc_count == 0

243 hc_dim = DIM // hc_count

244 assert block_output.shape == (N, hc_dim)

242 - assert injection_logits.shape == (N, hc_count)

245 assert residual.stride(1) == 1

246 assert block_output.stride(1) == 1

245 - assert injection_logits.stride(1) == 1

247 + if injection_logits is not None:

248 + assert injection_logits.shape == (N, hc_count)

249 + assert injection_logits.stride(1) == 1

250 +

251 + stride_injection = injection_logits.stride(0) if injection_logits is not None else 0

252

253 out = residual.new_empty(residual.shape)

254 BLOCK_SIZE = 512

@@ -253,7 +259,7 @@ def _hc_combine(

259 out,

260 block_output.stride(0),

261 residual.stride(0),

256 - injection_logits.stride(0),

262 + stride_injection,

263 out.stride(0),

264 hc_dim,

265 hc_count,

@@ -287,8 +293,9 @@ def _hc_combine_norm_kernel(

293 NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE)

294 NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)

295

290 - row = tl.program_id(0)

291 - stream = tl.program_id(1)

296 + pid = tl.program_id(0)

297 + row = pid // HC

298 + stream = pid % HC

299 offs_hc = tl.arange(0, HC_PAD)

300 mask_hc = offs_hc < HC

301 tile_ids = tl.arange(0, NUM_TILES_PAD)

@@ -305,27 +312,34 @@ def _hc_combine_norm_kernel(

312 # Start the uncached residual load first, then issue the other combine

313 # loads before consuming any of them.

314 res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0)

308 - inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)

309 - block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)

310 - inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)

311 - inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0))

315 + if inj_ptr is not None:

316 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)

317 + block = tl.load(

318 + block_ptr + row * stride_block + offs_inner,

319 + mask_inner,

320 + other=0.0,

321 + )

322 + if inj_ptr is not None:

323 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)

324 + block = block.to(tl.float32) * tl.sum(tl.where(offs_hc == stream, inj, 0.0))

325 # Round the materialized combine result before normalization. This matches

326 # the unfused combine -> RMSNorm boundary.

314 - out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty)

327 + out = (res.to(tl.float32) + block.to(tl.float32)).to(out_ptr.dtype.element_ty)

328 + if inj_ptr is None:

329 + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)

330 tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner)

331

332 out = out.to(tl.float32)

318 - # Keep the two-axis reduction: flattening the padded tile is ~40% slower

319 - # at decode sizes.

333 sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0)

334 rrms = tl.rsqrt(sum_sq / HC_DIM + EPS)

335

336 if launch_pdl:

337 tl.extra.cuda.gdc_launch_dependents()

338

326 - # Loading the weight earlier helps decode but keeps the tile live across

327 - # the reduction and regresses larger batches, so defer it to the norm.

328 - w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)

339 + if inj_ptr is not None:

340 + # Loading the weight earlier helps decode but keeps the tile live across

341 + # the reduction and regresses larger batches, so defer it to the norm.

342 + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)

343 y = out * rrms

344 y += y * w.to(tl.float32)

345 tl.store(y_ptr + row * stride_y + offs, y, mask_inner)

@@ -334,7 +348,7 @@ def _hc_combine_norm_kernel(

348 def _hc_combine_norm(

349 residual: torch.Tensor,

350 block_output: torch.Tensor,

337 - injection_logits: torch.Tensor,

351 + injection_logits: torch.Tensor | None,

352 norm_weight: torch.Tensor,

353 eps: float,

354 hc_count: int,

@@ -343,17 +357,20 @@ def _hc_combine_norm(

357 assert DIM % hc_count == 0

358 hc_dim = DIM // hc_count

359 assert block_output.shape == (N, hc_dim)

346 - assert injection_logits.shape == (N, hc_count)

360 assert residual.stride(1) == 1

361 assert block_output.stride(1) == 1

349 - assert injection_logits.stride(1) == 1

362 + if injection_logits is not None:

363 + assert injection_logits.shape == (N, hc_count)

364 + assert injection_logits.stride(1) == 1

365 assert norm_weight.is_contiguous()

366 assert norm_weight.numel() in (hc_dim, DIM)

367

368 + stride_injection = injection_logits.stride(0) if injection_logits is not None else 0

369 +

370 out = residual.new_empty(residual.shape)

371 y = residual.new_empty(residual.shape)

372 BLOCK_SIZE = 512

356 - _hc_combine_norm_kernel[(N, hc_count)](

373 + _hc_combine_norm_kernel[(N * hc_count,)](

374 block_output,

375 residual,

376 injection_logits,

@@ -362,7 +379,7 @@ def _hc_combine_norm(

379 y,

380 block_output.stride(0),

381 residual.stride(0),

365 - injection_logits.stride(0),

382 + stride_injection,

383 out.stride(0),

384 y.stride(0),

385 hc_dim,

@@ -389,7 +406,7 @@ def _hc_gate_mix_fake(

406 def _hc_combine_fake(

407 residual: torch.Tensor,

408 block_output: torch.Tensor,

392 - injection_logits: torch.Tensor,

409 + injection_logits: torch.Tensor | None,

410 hc_count: int,

411 ) -> torch.Tensor:

412 del block_output, injection_logits, hc_count

@@ -399,7 +416,7 @@ def _hc_combine_fake(

416 def _hc_combine_norm_fake(

417 residual: torch.Tensor,

418 block_output: torch.Tensor,

402 - injection_logits: torch.Tensor,

419 + injection_logits: torch.Tensor | None,

420 norm_weight: torch.Tensor,

421 eps: float,

422 hc_count: int,

@@ -452,7 +469,7 @@ def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Ten

469 def hc_combine(

470 residual: torch.Tensor,

471 block_output: torch.Tensor,

455 - injection_logits: torch.Tensor,

472 + injection_logits: torch.Tensor | None,

473 hc_count: int,

474 ) -> torch.Tensor:

475 return torch.ops.vllm.qwen4_exp_hc_combine(

@@ -463,7 +480,7 @@ def hc_combine(

480 def hc_combine_norm(

481 residual: torch.Tensor,

482 block_output: torch.Tensor,

466 - injection_logits: torch.Tensor,

483 + injection_logits: torch.Tensor | None,

484 norm_weight: torch.Tensor,

485 eps: float,

486 hc_count: int,