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

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 13, 17:00 UTC

[PERF] Decouple projections from GDN custom op. Attempt 2 (#28083)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>

Vadim Gimpelson committed Nov 6, 2025 at 05:01 UTC b6a248bdd7bff91a15c4f2395762f8ce4c19c759

3 files changed +205 -53

vllm/config/compilation.py

+1 -1

@@ -462,7 +462,7 @@ class CompilationConfig:

462 "vllm::short_conv",

463 "vllm::linear_attention",

464 "vllm::plamo2_mamba_mixer",

465 - "vllm::gdn_attention",

465 + "vllm::gdn_attention_core",

466 "vllm::kda_attention",

467 "vllm::sparse_attn_indexer",

468 ]

vllm/model_executor/layers/layernorm.py

+103

@@ -369,6 +369,109 @@ class GemmaRMSNorm(CustomOp):

369 return self.forward_native(x, residual)

370

371

372 +@CustomOp.register("rms_norm_gated")

373 +class RMSNormGated(CustomOp):

374 + """RMS Normalization with optional gating.

375 +

376 + This is a native PyTorch implementation that supports:

377 + - Standard RMS normalization

378 + - Group RMS normalization

379 + - Optional gating with SiLU activation

380 + """

381 +

382 + def __init__(

383 + self,

384 + hidden_size: int,

385 + eps: float = 1e-5,

386 + group_size: int | None = None,

387 + norm_before_gate: bool = False,

388 + device: torch.device | None = None,

389 + dtype: torch.dtype | None = None,

390 + ):

391 + """Initialize RMSNormGated.

392 +

393 + Args:

394 + hidden_size: Size of the hidden dimension

395 + eps: Epsilon for numerical stability

396 + group_size: If not None, do GroupNorm with each group

397 + having group_size elements.

398 + group_size=None is equivalent to group_size=hidden_size

399 + (i.e. there's only 1 group).

400 + norm_before_gate: If True and z is provided: out = norm(x) * silu(z)

401 + If False and z is provided: out = norm(x * silu(z))

402 + device: Device to create parameters on

403 + dtype: Data type for parameters

404 + """

405 + factory_kwargs = {"device": device, "dtype": dtype}

406 + super().__init__()

407 + self.eps = eps

408 + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

409 + self.register_parameter("bias", None)

410 + self.group_size = group_size

411 + self.norm_before_gate = norm_before_gate

412 + self.reset_parameters()

413 +

414 + def reset_parameters(self):

415 + torch.nn.init.ones_(self.weight)

416 +

417 + def forward_native(

418 + self, x: torch.Tensor, z: torch.Tensor | None = None

419 + ) -> torch.Tensor:

420 + """

421 + Native PyTorch implementation of RMS normalization with gating.

422 +

423 + Args:

424 + x: Input tensor

425 + z: Optional gating tensor

426 +

427 + Returns:

428 + Normalized (and optionally gated) tensor

429 +

430 + If z is not None:

431 + - norm_before_gate=True: out = norm(x) * silu(z)

432 + - norm_before_gate=False: out = norm(x * silu(z))

433 + """

434 + # Apply gating before normalization if needed

435 + if z is not None and not self.norm_before_gate:

436 + x = x * F.silu(z)

437 +

438 + # RMS Normalization

439 + if self.group_size is None:

440 + # Standard RMS norm across the last dimension

441 + variance = x.pow(2).mean(dim=-1, keepdim=True)

442 + x_normed = x * torch.rsqrt(variance + self.eps)

443 + out = x_normed * self.weight

444 + else:

445 + # Group RMS norm

446 + from einops import rearrange

447 +

448 + x_group = rearrange(x, "... (g d) -> ... g d", d=self.group_size)

449 + variance = x_group.pow(2).mean(dim=-1, keepdim=True)

450 + x_normed = x_group * torch.rsqrt(variance + self.eps)

451 + out = rearrange(x_normed, "... g d -> ... (g d)") * self.weight

452 +

453 + # Apply gating after normalization if needed

454 + if z is not None and self.norm_before_gate:

455 + out = out * F.silu(z)

456 +

457 + return out

458 +

459 + def forward_cuda(

460 + self, x: torch.Tensor, z: torch.Tensor | None = None

461 + ) -> torch.Tensor:

462 + from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn

463 +

464 + return rmsnorm_fn(

465 + x,

466 + self.weight,

467 + self.bias,

468 + z=z,

469 + eps=self.eps,

470 + group_size=self.group_size,

471 + norm_before_gate=self.norm_before_gate,

472 + )

473 +

474 +

475 class LayerNorm(nn.Module):

476 """

477 Layer Normalization.

vllm/model_executor/models/qwen3_next.py

+101 -52

@@ -30,12 +30,14 @@ from vllm.distributed import (

30 from vllm.forward_context import ForwardContext, get_forward_context

31 from vllm.logger import init_logger

32 from vllm.model_executor.layers.fla.ops import (

33 - RMSNormGated,

33 chunk_gated_delta_rule,

34 fused_recurrent_gated_delta_rule,

35 )

36 from vllm.model_executor.layers.fused_moe import SharedFusedMoE

38 -from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3NextRMSNorm

37 +from vllm.model_executor.layers.layernorm import (

38 + GemmaRMSNorm as Qwen3NextRMSNorm,

39 +)

40 +from vllm.model_executor.layers.layernorm import RMSNormGated

41 from vllm.model_executor.layers.linear import (

42 ColumnParallelLinear,

43 QKVParallelLinear,

@@ -436,17 +438,66 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

438 hidden_states: torch.Tensor,

439 output: torch.Tensor,

440 ):

439 - return torch.ops.vllm.gdn_attention(

440 - hidden_states,

441 - output,

441 + """

442 + Forward pass with three parts:

443 + 1. Input projection

444 + 2. Core attention (custom op)

445 + 3. Output projection

446 + """

447 + num_tokens = hidden_states.size(0)

448 +

449 + # ============================================================

450 + # Part 1: Input Projection

451 + # ============================================================

452 + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)

453 + projected_states_ba, _ = self.in_proj_ba(hidden_states)

454 + query, key, value, z, b, a = self.fix_query_key_value_ordering(

455 + projected_states_qkvz, projected_states_ba

456 + )

457 + query, key, value = map(

458 + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)

459 + )

460 + mixed_qkv = torch.cat((query, key, value), dim=-1)

461 +

462 + # ============================================================

463 + # Part 2: Core Attention (Custom Op)

464 + # ============================================================

465 + core_attn_out = torch.zeros(

466 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),

467 + dtype=hidden_states.dtype,

468 + device=hidden_states.device,

469 + )

470 +

471 + torch.ops.vllm.gdn_attention_core(

472 + mixed_qkv,

473 + b,

474 + a,

475 + core_attn_out,

476 self.prefix,

477 )

478

445 - def _forward(

479 + # ============================================================

480 + # Part 3: Output Projection

481 + # ============================================================

482 + z_shape_og = z.shape

483 + # Reshape input data into 2D tensor

484 + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

485 + z = z.reshape(-1, z.shape[-1])

486 + core_attn_out = self.norm(core_attn_out, z)

487 + core_attn_out = core_attn_out.reshape(z_shape_og)

488 + core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")

489 + output[:num_tokens], _ = self.out_proj(core_attn_out)

490 +

491 + def _forward_core(

492 self,

447 - hidden_states: torch.Tensor,

448 - output: torch.Tensor,

493 + mixed_qkv: torch.Tensor,

494 + b: torch.Tensor,

495 + a: torch.Tensor,

496 + core_attn_out: torch.Tensor,

497 ):

498 + """

499 + Core attention computation (called by custom op).

500 + """

501 forward_context = get_forward_context()

502 attn_metadata: AttentionMetadata = forward_context.attn_metadata

503

@@ -471,18 +522,11 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

522 num_actual_tokens = attn_metadata.num_actual_tokens

523 num_accepted_tokens = attn_metadata.num_accepted_tokens

524

474 - # 1. Set up dimensions for reshapes later

475 - projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states[:num_actual_tokens])

476 - projected_states_ba, _ = self.in_proj_ba(hidden_states[:num_actual_tokens])

477 - query, key, value, z, b, a = self.fix_query_key_value_ordering(

478 - projected_states_qkvz, projected_states_ba

479 - )

480 - query, key, value = map(

481 - lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)

482 - )

483 - mixed_qkv = torch.cat((query, key, value), dim=-1)

525 + mixed_qkv = mixed_qkv[:num_actual_tokens]

526 + b = b[:num_actual_tokens]

527 + a = a[:num_actual_tokens]

528

485 - # 2. Convolution sequence transformation

529 + # 1. Convolution sequence transformation

530 conv_weights = self.conv1d.weight.view(

531 self.conv1d.weight.size(0), self.conv1d.weight.size(2)

532 )

@@ -498,7 +542,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

542 mixed_qkv_spec = None

543 mixed_qkv_non_spec = mixed_qkv

544

501 - # 2.1: process the mutli-query part

545 + # 1.1: Process the multi-query part

546 if spec_sequence_masks is not None:

547 mixed_qkv_spec = causal_conv1d_update(

548 mixed_qkv_spec,

@@ -515,7 +559,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

559 validate_data=False,

560 )

561

518 - # 2.2: process the remaining part

562 + # 1.2: Process the remaining part

563 if attn_metadata.num_prefills > 0:

564 mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)

565 # - "cache_indices" updates the conv_state cache in positions

@@ -570,9 +614,9 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

614 g_non_spec = g

615 beta_non_spec = beta

616

573 - # 3. Recurrent attention

617 + # 2. Recurrent attention

618

575 - # 3.1: process the mutlti-query part

619 + # 2.1: Process the multi-query part

620 if spec_sequence_masks is not None:

621 core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(

622 q=query_spec,

@@ -590,7 +634,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

634 else:

635 core_attn_out_spec, last_recurrent_state = None, None

636

593 - # 3.2: process the remaining part

637 + # 2.2: Process the remaining part

638 if attn_metadata.num_prefills > 0:

639 initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()

640 initial_state[~has_initial_state, ...] = 0

@@ -633,30 +677,20 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

677 else:

678 core_attn_out_non_spec, last_recurrent_state = None, None

679

636 - # Merge core attention output

680 + # 3. Merge core attention output

681 if spec_sequence_masks is not None and core_attn_out_non_spec is not None:

638 - core_attn_out = torch.empty(

682 + merged_out = torch.empty(

683 (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),

684 dtype=core_attn_out_non_spec.dtype,

685 device=core_attn_out_non_spec.device,

686 )

643 - core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

644 - core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)

645 -

687 + merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

688 + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)

689 + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)

690 elif spec_sequence_masks is not None:

647 - core_attn_out = core_attn_out_spec

691 + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)

692 else:

649 - core_attn_out = core_attn_out_non_spec

650 -

651 - z_shape_og = z.shape

652 - # reshape input data into 2D tensor

653 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

654 - z = z.reshape(-1, z.shape[-1])

655 - core_attn_out = self.norm(core_attn_out, z)

656 - core_attn_out = core_attn_out.reshape(z_shape_og)

657 - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")

658 -

659 - output[:num_actual_tokens], _ = self.out_proj(core_attn_out)

693 + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)

694

695

696 class Qwen3NextAttention(nn.Module):

@@ -1260,29 +1294,44 @@ class Qwen3NextForCausalLM(

1294 return self.model.get_expert_mapping()

1295

1296

1263 -def gdn_attention(

1264 - hidden_states: torch.Tensor,

1265 - output: torch.Tensor,

1297 +def gdn_attention_core(

1298 + mixed_qkv: torch.Tensor,

1299 + b: torch.Tensor,

1300 + a: torch.Tensor,

1301 + core_attn_out: torch.Tensor,

1302 layer_name: str,

1303 ) -> None:

1304 + """

1305 + Custom op for the core attention computation.

1306 + Only handles the convolution + recurrent attention part.

1307 + Input/output projections are handled outside this op.

1308 + """

1309 forward_context: ForwardContext = get_forward_context()

1310 self = forward_context.no_compile_layers[layer_name]

1270 - self._forward(hidden_states=hidden_states, output=output)

1311 + self._forward_core(

1312 + mixed_qkv=mixed_qkv,

1313 + b=b,

1314 + a=a,

1315 + core_attn_out=core_attn_out,

1316 + )

1317

1318

1273 -def gdn_attention_fake(

1274 - hidden_states: torch.Tensor,

1275 - output: torch.Tensor,

1319 +def gdn_attention_core_fake(

1320 + mixed_qkv: torch.Tensor,

1321 + b: torch.Tensor,

1322 + a: torch.Tensor,

1323 + core_attn_out: torch.Tensor,

1324 layer_name: str,

1325 ) -> None:

1326 + """Fake implementation for torch.compile."""

1327 return

1328

1329

1330 direct_register_custom_op(

1282 - op_name="gdn_attention",

1283 - op_func=gdn_attention,

1284 - mutates_args=["output"],

1285 - fake_impl=gdn_attention_fake,

1331 + op_name="gdn_attention_core",

1332 + op_func=gdn_attention_core,

1333 + mutates_args=["core_attn_out"],

1334 + fake_impl=gdn_attention_core_fake,

1335 )

1336

1337