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

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

[Mamba] Add stochastic rounding support (#35753)

Signed-off-by: Roi Koren <roik@nvidia.com>

roikoren755 committed Mar 30, 2026 at 19:33 UTC 8e6293e838f94430d7aef14cbb7308c10b99ea11

7 files changed +166 -2

tests/kernels/mamba/test_mamba_ssm.py

+54

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (

12 selective_scan_fn,

13 selective_state_update,

14 )

15 +from vllm.platforms import current_platform

16 from vllm.utils.torch_utils import set_random_seed

17 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

@@ -429,6 +430,59 @@ def test_selective_state_update(dim, dstate, has_z, itype):

430 assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

431

432

433 +@pytest.mark.parametrize("philox_rounds", [0, 4])

434 +@pytest.mark.parametrize("has_z", [False, True])

435 +@pytest.mark.parametrize("dstate", [16, 64])

436 +@pytest.mark.parametrize("dim", [2048, 4096])

437 +@pytest.mark.skipif(

438 + not (

439 + current_platform.is_cuda() and current_platform.is_device_capability_family(100)

440 + ),

441 + reason="Stochastic rounding in triton is only supported"

442 + " on compute capability 10.0 CUDA devices.",

443 +)

444 +def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_rounds):

445 + device = "cuda"

446 + rtol, atol = 5e-3, 1e-1

447 + # set seed

448 + set_random_seed(0)

449 + batch_size = 1

450 + state = torch.randn(batch_size, dim, dstate, dtype=torch.float16, device=device)

451 + x = torch.randn(batch_size, dim, device=device, dtype=torch.bfloat16)

452 + out = torch.empty_like(x)

453 + dt = torch.randn(batch_size, dim, device=device, dtype=torch.bfloat16)

454 + dt_bias = torch.rand(dim, device=device) - 4.0

455 + A = -torch.rand(dim, dstate, device=device) - 1.0

456 + B = torch.randn(batch_size, dstate, device=device)

457 + C = torch.randn(batch_size, dstate, device=device)

458 + D = torch.randn(dim, device=device)

459 + z = torch.randn_like(x) if has_z else None

460 + # Reference uses fp32 state to get ground truth

461 + state_ref = state.float()

462 + selective_state_update(

463 + state,

464 + x,

465 + dt,

466 + A,

467 + B,

468 + C,

469 + D=D,

470 + z=z,

471 + dt_bias=dt_bias,

472 + dt_softplus=True,

473 + out=out,

474 + enable_stochastic_rounding=True,

475 + cache_philox_rounds=philox_rounds,

476 + )

477 + out_ref = selective_state_update_ref(

478 + state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True

479 + )

480 +

481 + assert state.dtype == torch.float16

482 + assert torch.allclose(state, state_ref.to(torch.float16), rtol=rtol, atol=atol)

483 + assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

484 +

485 +

486 @pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])

487 @pytest.mark.parametrize("has_z", [False, True])

488 @pytest.mark.parametrize("dstate", [16, 64])

vllm/config/cache.py

+35 -1

@@ -109,12 +109,20 @@ class CacheConfig:

109 mamba_cache_mode: MambaCacheMode = "none"

110 """The cache strategy for Mamba layers.

111 - "none": set when prefix caching is disabled.

112 - - "all": cache the mamba state of all tokens at position i * block_size. This is

112 + - "all": cache the mamba state of all tokens at position i * block_size. This is

113 the default behavior (for models that support it) when prefix caching is

114 enabled.

115 - "align": only cache the mamba state of the last token of each scheduler step and

116 when the token is at position i * block_size.

117 """

118 + enable_mamba_cache_stochastic_rounding: bool = False

119 + """Enable stochastic rounding when writing SSM state to fp16 cache.

120 + Uses random bits to unbias the rounding error, which can improve

121 + numerical stability for long sequences."""

122 + mamba_cache_philox_rounds: int = 0

123 + """Number of Philox PRNG rounds for stochastic rounding random number

124 + generation. 0 uses the Triton default. Higher values improve randomness

125 + quality at the cost of compute."""

126

127 # Will be set after profiling.

128 num_gpu_blocks: int | None = field(default=None, init=False)

@@ -231,3 +239,29 @@ class CacheConfig:

239 "scaling factor."

240 )

241 return cache_dtype

242 +

243 + def __post_init__(self):

244 + if self.enable_mamba_cache_stochastic_rounding:

245 + from vllm.platforms import current_platform

246 +

247 + if not current_platform.is_cuda():

248 + raise ValueError(

249 + "Stochastic rounding for Mamba cache is only supported "

250 + "on NVIDIA CUDA platforms. Please do not specify "

251 + "`--enable-mamba-cache-stochastic-rounding`."

252 + )

253 + if not current_platform.is_device_capability_family(100):

254 + raise ValueError(

255 + "Stochastic rounding for Mamba cache requires compute "

256 + "capability 10.0 (data center Blackwell). The `cvt.rs` PTX "

257 + "instruction is not supported on your GPU. Please do not specify "

258 + "`--enable-mamba-cache-stochastic-rounding`."

259 + )

260 + if self.mamba_ssm_cache_dtype != "float16":

261 + raise ValueError(

262 + "Stochastic rounding for Mamba cache requires "

263 + "the SSM cache to be float16. Please set it explicitly, "

264 + "by specifying `--mamba-ssm-cache-dtype float16`, or disable "

265 + "stochastic rounding by not specifying "

266 + "`--enable-mamba-cache-stochastic-rounding`."

267 + )

vllm/engine/arg_utils.py

+13

@@ -604,6 +604,10 @@ class EngineArgs:

604 mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype

605 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")

606 mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode

607 + enable_mamba_cache_stochastic_rounding: bool = (

608 + CacheConfig.enable_mamba_cache_stochastic_rounding

609 + )

610 + mamba_cache_philox_rounds: int = CacheConfig.mamba_cache_philox_rounds

611

612 additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")

613

@@ -1024,6 +1028,13 @@ class EngineArgs:

1028 cache_group.add_argument(

1029 "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]

1030 )

1031 + cache_group.add_argument(

1032 + "--enable-mamba-cache-stochastic-rounding",

1033 + **cache_kwargs["enable_mamba_cache_stochastic_rounding"],

1034 + )

1035 + cache_group.add_argument(

1036 + "--mamba-cache-philox-rounds", **cache_kwargs["mamba_cache_philox_rounds"]

1037 + )

1038 cache_group.add_argument(

1039 "--kv-offloading-size", **cache_kwargs["kv_offloading_size"]

1040 )

@@ -1590,6 +1601,8 @@ class EngineArgs:

1601 mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,

1602 mamba_block_size=self.mamba_block_size,

1603 mamba_cache_mode=self.mamba_cache_mode,

1604 + enable_mamba_cache_stochastic_rounding=self.enable_mamba_cache_stochastic_rounding,

1605 + mamba_cache_philox_rounds=self.mamba_cache_philox_rounds,

1606 kv_offloading_size=self.kv_offloading_size,

1607 kv_offloading_backend=self.kv_offloading_backend,

1608 )

vllm/model_executor/layers/mamba/mamba_mixer.py

@@ -428,6 +428,8 @@ class MambaMixer(MambaBase, PluggableLayer):

428 state_batch_indices=state_indices_tensor_d_input,

429 dst_state_batch_indices=state_indices_tensor_d_output,

430 out=scan_outputs_d,

431 + enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,

432 + cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,

433 )

434 scan_outputs_d = scan_outputs_d.transpose(0, 1)

435

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -888,6 +888,8 @@ class MambaMixer2(MambaBase, PluggableLayer):

888 num_accepted_tokens=num_accepted_tokens,

889 cu_seqlens=query_start_loc_d,

890 is_blackwell=self.is_blackwell,

891 + enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,

892 + cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,

893 )

894

895 def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

+58 -1

@@ -28,6 +28,21 @@ else:

28 return dt

31 +@triton.jit

32 +def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor:

33 + y = tl.inline_asm_elementwise(

34 + asm="""{

35 +cvt.rs.f16x2.f32 $0, $2, $1, $3;

36 +}""",

37 + constraints="=r,r,r,r,r",

38 + args=(x, rand),

39 + dtype=tl.float16,

40 + is_pure=True,

41 + pack=2,

42 + )

43 + return y

44 +

45 +

46 @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})

47 @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})

48 @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})

@@ -48,6 +63,7 @@ else:

63 def _selective_scan_update_kernel(

64 # Pointers to matrices

65 state_ptr,

66 + rand_seed_ptr,

67 x_ptr,

68 dt_ptr,

69 dt_bias_ptr,

@@ -113,6 +129,8 @@ def _selective_scan_update_kernel(

129 IS_SPEC_DECODING: tl.constexpr,

130 IS_VARLEN: tl.constexpr,

131 BLOCK_SIZE_DSTATE: tl.constexpr,

132 + USE_RS_ROUNDING: tl.constexpr,

133 + PHILOX_ROUNDS: tl.constexpr,

134 ):

135 pid_m = tl.program_id(axis=0)

136 pid_b = tl.program_id(axis=1)

@@ -267,7 +285,35 @@ def _selective_scan_update_kernel(

285 z_ptr += stride_z_batch

286

287 if not IS_SPEC_DECODING:

270 - tl.store(dst_state_ptrs, state.to(dst_state_ptrs.dtype.element_ty), mask=mask)

288 + if USE_RS_ROUNDING:

289 + # Load random seed

290 + rand_seed = tl.load(rand_seed_ptr)

291 + # Generate random offsets for each element in state

292 + if HAS_STATE_BATCH_INDICES:

293 + rand_offsets = (

294 + state_batch_idx * stride_state_batch + pid_h * stride_state_head

295 + )

296 + else:

297 + rand_offsets = pid_b * stride_state_batch + pid_h * stride_state_head

298 + rand_offsets += (

299 + offs_m[:, None] * stride_state_dim

300 + + offs_n[None, :] * stride_state_dstate

301 + )

302 + # Generate random 32-bits for each element in state

303 + if PHILOX_ROUNDS > 0:

304 + rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)

305 + else:

306 + rand = tl.randint(rand_seed, rand_offsets)

307 + # Convert state to fp16 with RS rounding

308 + state = convert_rs_fp16x2(state, rand)

309 + tl.static_assert(state.dtype == tl.float16, "state must be fp16")

310 + tl.static_assert(

311 + dst_state_ptrs.dtype.element_ty == tl.float16,

312 + "dst_state_ptrs must be fp16",

313 + )

314 + else:

315 + state = state.to(dst_state_ptrs.dtype.element_ty)

316 + tl.store(dst_state_ptrs, state, mask=mask)

317

318

319 def selective_state_update(

@@ -288,6 +334,8 @@ def selective_state_update(

334 num_accepted_tokens=None,

335 cu_seqlens=None,

336 is_blackwell=False,

337 + enable_stochastic_rounding=False,

338 + cache_philox_rounds=0,

339 ):

340 """

341 Argument:

@@ -419,9 +467,16 @@ def selective_state_update(

467 and dt.stride(-1) == 0

468 and dt_bias.stride(-1) == 0

469 )

470 + rand_seed = (

471 + torch.randint(0, 2**32, (1,), device=state.device)

472 + if enable_stochastic_rounding

473 + else None

474 + )

475 +

476 with torch.accelerator.device_index(x.device.index):

477 _selective_scan_update_kernel[grid](

478 state,

479 + rand_seed,

480 x,

481 dt,

482 dt_bias,

@@ -476,6 +531,8 @@ def selective_state_update(

531 tie_hdim,

532 BLOCK_SIZE_M,

533 num_warps=num_warps,

534 + USE_RS_ROUNDING=enable_stochastic_rounding,

535 + PHILOX_ROUNDS=cache_philox_rounds,

536 )

537

538

vllm/model_executor/models/plamo2.py

@@ -445,6 +445,8 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):

445 dt_softplus=True,

446 state_batch_indices=state_indices_tensor_d,

447 out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),

448 + enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,

449 + cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,

450 )

451

452 # 4. Final linear projection