# One-Dimensional Spin-Polarised Surface States - A Comparison of Bi(112) with
Other Vicinal Bismuth Surfaces

Anna Cecilie Åsland  Johannes Bakkelund  Department of Physics, Norwegian
University of Science and Technology (NTNU), NO-7491 Trondheim, Norway.  Even
Thingstad  Department of Physics, Norwegian University of Science and
Technology (NTNU), NO-7491 Trondheim, Norway.  Department of Physics,
University of Basel, Klingelbergstrasse 82, CH-4056 Basel, Switzerland.  Håkon
I. Røst  Department of Physics, Norwegian University of Science and Technology
(NTNU), NO-7491 Trondheim, Norway.  Simon P. Cooil  Centre for Materials
Science and Nanotechnology, University of Oslo, Oslo 0318, Norway.  Jinbang Hu
Department of Physics, Norwegian University of Science and Technology (NTNU),
NO-7491 Trondheim, Norway.  Ivana Vobornik  Jun Fujii  Istituto Officina dei
Materiali (IOM)-CNR, Area Science Park-Basovizza, S.S. 14 Km 163.5, 34149
Trieste, Italy.  Asle Sudbø  Department of Physics, Norwegian University of
Science and Technology (NTNU), NO-7491 Trondheim, Norway.  Justin W. Wells
<EMAIL_ADDRESS>Department of Physics, Norwegian University of Science
and Technology (NTNU), NO-7491 Trondheim, Norway.  Centre for Materials
Science and Nanotechnology, University of Oslo, Oslo 0318, Norway.  Federico
Mazzola  Istituto Officina dei Materiali (IOM)-CNR, Area Science Park-
Basovizza, S.S. 14 Km 163.5, 34149 Trieste, Italy.

###### Abstract

Vicinal surfaces of bismuth are unique test-beds for investigating one-
dimensional (1D) spin-polarised surface states that may one day be used in
spintronic devices. In this work, two such states have been observed for the
(112) surface when measured using angle-resolved photoemission spectroscopy
(ARPES) and spin-resolved ARPES, and when calculated using a tight-binding
(TB) model and with density functional theory (DFT). The surface states appear
as elongated Dirac-cones which are 1D and almost dispersionless in the
$\mathbf{k}_{\text{y}}$-direction, but disperse with energy in the orthogonal
$\mathbf{k}_{\text{x}}$-direction to form two “$\times$”-like features
centered at the $\mathbf{k}_{\text{y}}$-line through $\bar{\Gamma}$. Unlike
many materials considered for spintronic applications, their 1D nature
suggests that conductivity and spin-transport properties are highly dependent
on direction. The spin-polarisation of the surface states is mainly in-plane
and parallel to the 1D state, but there are signs of a tilted out-of-plane
spin-component for one of them. The Bi(112) surface states resemble those
found for other vicinal surfaces of bismuth, strongly indicating that their
existence and general properties are robust properties of vicinal surfaces of
bismuth. Furthermore, differences in the details of the states, particularly
related to their spin-polarisation, suggest that spin-transport properties may
be engineered simply by precise cutting and polishing of the crystal.

††preprint: APS/123-QED

##  I Introduction

To find materials where spin-filtering and spin-transport can be realised,
either through electrical currents or magnon currents, are of major importance
in spintronics [1, 2, 3, 4, 5, 6]. Spin states facilitating the flow of spin-
polarised electrical currents are typically found in heavy element materials
like bismuth, because of their strong spin-orbit coupling [7, 8, 9, 10].

Two surfaces of bismuth, Bi(441) and Bi(114), have previously been found to
support one-dimensional (1D), spin-polarised surface states [11, 12]. The one-
dimensionality of these states suggests that conductivity and spin-transport
can vary along different directions on the surface. This distinguishes them
from many topological insulators for which conductivity and spin-transport
properties tend to be closer to isotropic [13, 14, 15, 16]. Furthermore, the
spin-polarisation of the surface states on Bi(441) and Bi(114) makes back-
scattering of surface state electrons less likely since this implies a full
reversal of their spin vectors [3, 4, 5]. Another advantage of pure bismuth is
how easy it is to cut, polish and clean compared to many topological
insulators, which tend to be alloys [13, 14, 15, 16]. True vicinal surfaces of
alloys are challenging to achieve consistently, as the surface stoichiometry
and therefore the electronic properties will easily vary upon re-preparation.

One of the reasons for the intense interest in topological insulators is that
the topology of the surface states protect electrons in these states from
scattering, unless the scattering involves a mechanism that allows exchange of
spin angular momentum [13, 3, 4, 5, 14, 15, 16]. Since topological insulators
are ideally insulators in the bulk, electrical transport happens through their
surface states. Bismuth is a semi-metal, so transport can happen through bulk
states, where electrons are not generally protected from scattering, in
addition to through surface states [17, 8]. However, if gaps could be opened
in the energy spectrum of bismuth such that only surface states contribute to
the Fermi-surface, or alternatively if vicinal surfaces could be achieved on
similar topological insulators like $\mathrm{Bi}_{1-x}\mathrm{Sb}_{x}$ [18, 5,
19], they would be prime candidates for anisotropic electrical spin-transport.
Vicinal surfaces of bismuth are therefore great test-beds to investigate the
properties of such surface states, and understanding the underlying physical
mechanisms should aid the exploration of multi-functional spin-based devices
in the future.

Here we show that another vicinal surface of bismuth, Bi(112), exhibits 1D,
spin-polarised surface states resembling those found on Bi(441) and Bi(114)
[11, 12]. Angle-resolved photoemission spectroscopy (ARPES) measurements
reveal two spin-polarised “$\times$”-like features in the energy spectrum,
centered at $\bar{\Gamma}_{1}$ in the centre of the 1st Brillouin zone (BZ,
see Fig. 1(d)), with each of their crossing points at approximately the same
binding energies ($E_{\text{B}}$, also known as $E-E_{\text{F}}$) as the
surface states of Bi(441) [11]. Similarly, Bi(114) also has a 1D spin-
polarised surface state at one of these $E_{\text{B}}$ [12]. Spin-polarised
surface states therefore appear to be a robust property of vicinal bismuth
surfaces. Interestingly, all three surfaces show subtle differences in the
details of the surface states, suggesting that properties like the direction
of the spin-vector can be tuned simply by cutting the crystal along a given
direction. This knowledge can possibly be exploited to create easily
modifiable materials for use in new spintronic devices.

##  II Results & Discussion

Figure 1: The atomic and electronic structure of Bi(112). (a) Sketch showing
the atomic structure of the vicinal Bi(112) surface from the side. Red, green
and blue spheres represent the uppermost, 2nd and 3rd layers of the crystal
respectively. (b) Top view schematic of the Bi(112) surface with the unit cell
indicated (green). (c) Low-energy electron diffraction (LEED) pattern of the
Bi(112) surface, measured with electron energy 40 eV. (d) The two-dimensional
Fermi-surface of Bi(112). The 1st and 2nd Brillouin zones (BZ) are indicated
by green and purple rectangles. (e) Volumetric representation of the measured
Bi(112) band structure. (f) Measured band structure ($E$ vs
$\mathbf{k}_{\text{x}}$) at $\mathbf{k}_{\text{y}}=0$, i.e. orthogonal to the
one-dimensional surface states.

The vicinal surface of Bi(112) can be described by sheets of atoms with edges
forming parallel lines on the surface (see Fig. 1(a) and (b)). Due to
dimerisation, the lines are only weakly coupled. This results in a macroscopic
number of parallel quasi 1D systems. The surface structure shown in Figs. 1(a)
and (b) is confirmed by low-energy electron diffraction (LEED) in Fig. 1(c).
The distance between the streaks indicated in Fig. 1(c) matches the
periodicity of the truncated bulk crystal along the direction orthogonal to
the 1D atomic chains. ARPES measurements further confirm the 1D character of
the surface localised electrons. The Fermi-surface with 1D states along
$\mathbf{k}_{\text{y}}$ is shown in Fig. 1(d). As seen in Figs. 1(e) and (f),
these lines consist of two bands in each BZ dispersing in energy to form
“$\times$”-like features, crossing each other at the $\bar{\Gamma}$-points
close to the Fermi-level. The bands resemble the shape of a Dirac cone, but
with the Dirac point elongated to form a “Dirac line” parallel to
$\mathbf{k}_{\text{y}}$ and an “$\times$” in the orthogonal direction along
$\mathbf{k}_{\text{x}}$. Additional features in the measured spectrum are
identified as bulk bands by comparison to density functional theory (DFT)
calculations, see details in the Supplementary Material [20].

Figure 2: (a) Overview of the band structure near the $\bar{\Gamma}_{1}$-point
($E$ vs $\mathbf{k_{\text{x}}}$) at $\mathbf{k_{\text{y}}}=0$. Green and pink
bands are bulk and surface states, respectively, calculated using a tight-
binding (TB) model. Dark blue and light blue bands represent the two
“$\times$”-like features at $\bar{\Gamma}_{1}$ as seen in the measurements.
Arrows from $\times_{E_{\text{F}}}$ and $\times_{0.16\,\text{eV}}$ indicate
the crossing points of the $\times$’s. The purple circles and arrow show where
$\times_{E_{\text{F}}}$ cross $\times_{0.16\,\text{eV}}$ away from
$\bar{\Gamma}_{1}$ in the measurements (dark blue cross light blue) and in the
calculations (pink). (b)-(d), (f)-(h) Band structure along
$\bar{\mathrm{X}}-\bar{\Gamma}_{1}-\bar{\mathrm{X}}$ measured at
photoexcitation energies $h\nu=21$ eV (b), $h\nu=32$ eV (c)-(d), $45$ eV (f),
$70$ eV (g) and $80$ eV (h). TB calculations are overlaid on the measurements.
In (d), which is measured inside the orange region in (c), the left hand side
is measured with $p$-polarised light, and the right hand side with
$s$-polarised light. The direction of the light polarisation with respect to
the sample surface is shown in (e).

In order to differentiate surface states from bulk states, measurements were
performed as a function of photoexcitation energy $h\nu$. Although symmetry
breaking in the out-of-plane direction ($\hat{\mathbf{z}}$) caused by the
surface means that the out-of-plane momentum ($\propto\mathbf{k}_{\text{z}}$)
is not well conserved in the photoemission process, the
$\mathbf{k}_{\text{z}}$ dispersion of emitted bulk state electrons can still
be probed by varying the excitation energy [21]. On the other hand, surface
states are localised on the surface and their energy is therefore independent
of $\mathbf{k}_{\text{z}}$ (and $h\nu$) [22].

An overview of the results is given in Fig. 2(a), showing a sketch of the
measured surface states (dark blue and light blue), compared to calculated
tight-binding (TB) surface (pink) and bulk (green) states. In Figs. 2(b)-(d)
and (f)-(h), the bands at $E_{\text{B}}$ larger than $\approx 0.3$ eV are seen
to disperse with photoexcitation energy and are therefore identified as bulk
bands. In addition, there are two “$\times$”-like features formed by bands
crossing each other at $\bar{\Gamma}_{1}$, i.e. $\times_{E_{\text{F}}}$ and
$\times_{0.16\,\text{eV}}$, shown as dark blue and light blue bands
respectively in Fig. 2(a). These bands do not disperse with photoexcitation
energy, hence they are identified as surface states. However, a noticeable
variation in intensity is observed, most likely related to matrix element
effects [21, 23]. As indicated by arrows in Fig. 2(a), $\times_{E_{\text{F}}}$
(dark blue) has its crossing point near the Fermi-level at
$E_{\text{B}}=0.07\pm 0.10$ eV. Whereas, $\times_{0.16\,\text{eV}}$ (light
blue in Fig. 2(a)) has the crossing point at a higher $E_{\text{B}}$ of
$E_{\text{B}}=0.16\pm 0.05$ eV, and can be seen from the measurements in Fig.
2(c). See the Supplementary Material for details on how the $E_{\text{B}}$ of
the crossing points were determined [20].

The electronic structure of Bi(112) was calculated using TB and DFT. Results
from the DFT-calculations are included in the Supplementary Material [20]. The
TB calculations are underlying the $\times$’s in Fig. 2(a) and overlaid on the
measured band structures in Figs. 2(b) and (c). The bulk states (green) have
been calculated for a range of $\mathbf{k}_{\text{z}}$, hence the measured
bulk states are expected to disperse within the area they cover. There is good
agreement between the TB calculated (green) and measured bulk states.

The TB calculated surface states (pink) show band structure similar to the
measurement in terms of two $\times$-like features with crossing points at
$E_{\text{B}}=-0.032$ eV and $E_{\text{B}}=0.166$ eV for
$\times_{E_{\text{F}}}$ and $\times_{0.16\,\text{eV}}$, respectively. Although
the $E_{\text{B}}$ of the crossing points in the measurements agree within
uncertainties with the values expected from calculations, the gradients of the
bands differ. Below the $\times_{E_{\text{F}}}$ crossing point and above the
$\times_{0.16\,\text{eV}}$ crossing point, the measured bands are much steeper
than the calculated bands. Because of this, the point where the bands
extending from $\times_{E_{\text{F}}}$ and $\times_{0.16\,\text{eV}}$ are
measured to cross (purple circles and arrow in Fig. 2(a)) is at a different
value of $\mathbf{k}_{\text{x}}$ and $E_{\text{B}}$ compared to the calculated
bands. Preliminary calculations indicated that surface state gradients and
energies are very sensitive to the precise geometry at the surface [24].
Relaxation of the surface, or missing rows of atoms as observed in Ref. 12,
may therefore explain the difference in band-gradients, potentially leading to
a mechanism to control the surfance Fermi velocities. Such missing rows may
also explain why the surface states appear more 1D than expected from the
calculations: if the distance between the sheet edges is greater, a weaker
interaction between their constituent electronic states is expected [25].

[b]  Surface $E_{\text{B}}$ Crossing Point [eV] Spin-Vector Angle [∘]
$\times_{E_{\text{F}}}$ $\times_{0.16\,\text{eV}}$ $\times_{E_{\text{F}}}$
$\times_{0.16\,\text{eV}}$ Bi(112) $0.07\pm 0.10$ $0.16\pm 0.05$ $0$ $27\pm 7$
Bi(441) [11] $\approx 0$ $0.15-0.20$ $0$ $45$ Bi(114) [12] $0-0.15$a Not seen
$30$ Not seen TB Bi(112) $-0.032$ $0.166$ $\approx 0$ $\approx 0$

Table 1: Comparison of the binding energy ($E_{\text{B}}$) of the crossing
points and spin-vector of the $\times$-like features in the measured band
structure of three surfaces of bismuth. TB Bi(112) gives the equivalent
results from tight-binding calculations. The spin-vector angle gives the angle
between the spin-vector and the surface plane, i.e. a measure of how large the
out-of-plane spin-component is.

  * a

For Bi(114), the surface state looks like a single line feature on top of a
$\Lambda$-like bulk state, hence it is not clear whether the surface state is
$\times$-like with a crossing point. The maximum $E_{\text{B}}$ in the given
range is estimated as the highest $E_{\text{B}}$ of the surface state. See
details in Ref. 12.

The $\times$-like features are similar to the surface states observed for
other vicinal surfaces of bismuth, as summarised in Table 1. Two $\times$-like
surface states with their crossing point near the Fermi-level and at
$E_{\text{B}}=0.16-0.20$ eV are observed for both the (112) and (441) surface
[11]. The band structure of Bi(114) has a $\Lambda$-shaped bulk feature with
its maximum point near $E_{\text{B}}\approx 0.15$ eV, but with an additional
1D surface state at $E_{\text{B}}$ between $0-0.15$ eV [12]. For all three
surfaces, the one-dimensionality of the surface states is expected to
influence the conduction properties. A distinct increase in the conductivity
along the atomic rows of the surface when compared to the orthogonal direction
is expected as there are far more states to scatter into along the surface
chains (see Fig. 1(b)).

In order to investigate the orbital origins of the observed states, the
polarisation of the light was varied between $s$\- and $p$\- polarisation (in-
plane and partly out-of-plane respectively) as visualised in figure 2(e). An
example is given in Fig. 2(d), where an additional band can be observed inside
the $\times_{0.16\,\text{eV}}$-feature when measured with $s$-polarised light
(right hand side panel). Since this is in the region where bulk bands are
found from the calculations, it is believed to be a bulk band. See the
Supplementary Material for more details on the orbital nature of the states
[20].

Figure 3: Spin-polarisation from tight-binding calculations. (a), (b), (c)
Band structure between $\bar{\text{X}}-\bar{\Gamma}_{1}-\bar{\text{X}}$. In
(a) and (b) the colour-scale shows the spin-polarisation
$\mathbf{S}_{\text{y}}$, the spin-vector component in-plane and along the 1D
lines, (a), and $\mathbf{S}_{\text{z}}$, the out-of-plane component of the
spin-vector, (b). In (c), the colours label the surface bands. (d), (e) Spin-
polarisation of the spin-vector component $\mathbf{S}_{\text{y}}$ (d) and
$\mathbf{S}_{\text{z}}$ (e) respectively. The colour shows which band in (c)
the spin-polarisation value belongs to.

Due to the strong spin-orbit coupling in bismuth, the surface states are
expected to be spin-polarised [8, 7]. This is demonstrated in Fig. 3, showing
the TB calculated spin-polarisation dependence of the surface states. Positive
(negative) spin-polarisation shown in red (blue) gives the magnitude of the
spin-vector component in-plane and pointing along $+\mathbf{k}_{\text{y}}$
($-\mathbf{k}_{\text{y}}$), $\mathbf{S}_{\text{y}}$, in Fig. 3(a) and out-of-
plane pointing along $+\mathbf{k}_{\text{z}}$ ($-\mathbf{k}_{\text{z}}$),
$\mathbf{S}_{\text{z}}$, in Fig. 3(b). For all the surface bands,
$\mathbf{S}_{\text{y}}$ is close to $\pm 100$ % (Figs. 3(a) and (d)). The
exception is the points where the surface states corresponding to the two
crosses $\times_{0.16\,\text{eV}}$ and $\times_{E_{\text{F}}}$ cross (purple
circles and arrow in Fig. 2(a)). There, the surface states also attain a
significant out-of plane spin-component $\mathbf{S}_{\text{z}}$ (Figs. 3(b)
and (e)).

The measured in-plane spin-polarisation shown in Fig. 4 is consistent with the
calculations. In Figs. 4(a) and (b), constant energy surfaces are shown at
$E_{\text{B}}=0.12$ eV and $E_{\text{B}}=0.72$ eV. At both energies, the
region spanned by the 1st BZ (green, dashed rectangle) have been filled with
the measured projection of the spin-vector in-plane and along
$\mathbf{k}_{\text{y}}$ ($\mathbf{S}_{\text{y}}$). At $E_{\text{B}}=0.12$ eV,
the positive spin-polarisation (red) occurs for negative values of
$\mathbf{k}_{\text{x}}$, while at $E_{\text{B}}=0.72$ eV, the positive spin-
polarisation is at positive values of $\mathbf{k}_{\text{x}}$. This indicates
that the spin-polarisation reverses at a band crossing between
$E_{\text{B}}=0.12-0.72$ eV. The spin-reversal is confirmed in Figs. 4(c) and
(f) where the in-plane spin-polarisation $\mathbf{S}_{\text{y}}$ along the 1D
line is overlaid on the $E$ vs $\mathbf{k}_{\text{x}}$ maps of the surface
states. $\mathbf{S}_{\text{y}}$ was found to be $\approx 30$ % when measured
above and below the crossing point for $\times_{0.16\,\text{eV}}$, and
$\approx 70$ % for $\times_{E_{\text{F}}}$, see Figs. 4(d), (e) and (g),
respectively.

Figure 4: Measured spin texture. (a), (b) Constant energy surfaces at binding
energies $E_{\text{B}}=0.12$ eV (a) and $E_{\text{B}}=0.72$ eV (b). The first
Brillouin zone is substituted by the in-plane spin-polarisation along the
$\mathbf{k}_{\text{y}}$-direction visualised by red and blue arrows. (c), (f)
$E$ vs $\mathbf{k}_{\text{x}}$ plot measured with photoexcitation energies of
$h\nu=32$ eV and $h\nu=21$ eV. The in-plane spin-polarisation along
$\mathbf{k}_{\text{y}}$ was measured along the dashed white lines and has been
overlaid. (d), (e), (g) In-plane spin-polarisation along
$\mathbf{k}_{\text{y}}$ at $E_{\text{B}}=0$ eV (d) and $E_{\text{B}}=0.63$ eV
(e) in (c), and $E_{\text{B}}=0.36$ eV (g) in (f).

It should be noted that the spin-texture is similar for both
$\times_{E_{\text{F}}}$ and $\times_{0.16\,\text{eV}}$ although their
magnitude differ, as seen when comparing Figs. 4(e) and (g). Because of this,
we presume that backscattering is strongly reduced for electronic states in
the direction corresponding to the atomic chains in Fig. 1(b). In addition,
the spin-polarisation means that instabilities like charge density waves are
not expected on the Bi(112) surface [26, 27].

Signs of an out-of-plane spin-vector component can also be seen when measuring
the spin-polarisation of $\times_{0.16\,\text{eV}}$ (see details in the
Supplementary Material [20]). However, this spin-component is small compared
to $\mathbf{S}_{\text{y}}$, and approaches the detection limit in our
measurements. Even though $\mathbf{S}_{\text{z}}$ is small in the measurements
and is only significant for short $\mathbf{k}_{\text{x}}$-ranges in the
calculations, it can explain why the measured $\mathbf{S}_{\text{y}}$ in Fig.
4(d) is lower than expected from the calculations. These spin measurements
were performed near the Fermi-level, which is close to where
$\times_{0.16\,\text{eV}}$ and $\times_{E_{\text{F}}}$ cross (purple circles
and arrow in Fig. 2(a)). At this point, $\mathbf{S}_{\text{z}}$ increases
drastically in the calculations, and consequently $\mathbf{S}_{\text{y}}$
decreases.

When comparing the spin-polarisation found for Bi(112) with Bi(441) and
Bi(114) (see Table 1), they are found to be similar, but with important
differences [11, 12]. All three surfaces have 1D surface states where the in-
plane spin-component along the 1D states ($\mathbf{S}_{\text{y}}$) is the
strongest component. For $\times_{E_{\text{F}}}$ in Bi(112) and Bi(441) this
is the only component. $\times_{E_{\text{F}}}$ in Bi(114) and
$\times_{0.16\,\text{eV}}$ in Bi(441) was previously found to have an out-of-
plane component in addition to the in-plane component, such that the spin-
vector makes an angle $\approx 30^{\circ}$ and $\approx 45^{\circ}$ with the
(114) and (441) surface plane respectively [11, 12]. Bianchi et al. (2015)
[11] suggested that the out-of-plane component of the spin may be a
consequence of exposing edges of the (111) plane in bismuth (see Fig. 1(a)),
and that this spin-vector is in fact in-plane in the (111) plane. This would
mean that $\times_{E_{\text{F}}}$ in Bi(114) and $\times_{0.16\,\text{eV}}$ in
Bi(441) correspond to surface states in Bi(111), while $\times_{E_{\text{F}}}$
in Bi(441) corresponds to surface states on the (441) surface. This hypothesis
seems to match well for Bi(441) and Bi(114) [11, 12] when comparing the angle
of the spin-vector with the angle between the (111) and (114), or (441)
surface planes.

To see whether $\times_{E_{\text{F}}}$ and $\times_{0.16\,\text{eV}}$ in
Bi(112) also can be related to states on different surfaces, the full spin-
vector was estimated. As mentioned earlier, the spin-vector of
$\times_{E_{\text{F}}}$ seems to be fully in-plane, hence indicating that it
belongs to a surface state on Bi(112). $\times_{0.16\,\text{eV}}$ on the other
hand, seems to have a small out-of-plane component, causing the spin-vector to
be tilted out-of-plane to make an angle of $(27\pm 7)^{\circ}$ with the
surface plane. This angle is smaller than the angle of $37^{\circ}$ between
the (111) and (112) plane, but almost within the uncertainty of the estimate.
Altogether, the three surfaces have similar spin-vectors, but the angle
between the surface plane and spin-vector seems to vary. These subtle but
important differences mean that the spin-vector potentially can be tuned by
choosing a particular vicinal surface.

##  III Conclusion

One-dimensional, spin-polarised surface states have been observed on the (112)
surface of bismuth. The surface states are seen as “$\times$”-like features in
the band structure with crossing points at $E_{\text{B}}=0.07\pm 0.10$ eV and
$E_{\text{B}}=0.16\pm 0.05$ eV for $\times_{E_{\text{F}}}$ and
$\times_{0.16\,\text{eV}}$ respectively. These values agree within
uncertainties with calculations, even if the gradients of the measured bands
differ from the calculations. The spin-vector of the surface states is mainly
in-plane and along the 1D lines, but $\times_{0.16\,\text{eV}}$ shows
indications of having an additional out-of-plane spin-component, approximately
matching with the angle of the (111) plane. The surface states observed in
Bi(112) resemble those found for other vicinal surfaces of bismuth, indicating
that the existence of such states is a robust property of these surfaces.
Furthermore, variations in how the spin-vector is tilted out-of-plane for
different surfaces, can make it possible to tune the spin-transport on
bismuth-surfaces by adjusting the surface crystallographic direction.

##  IV Acknowledgements

This work was partly supported by the Research Council of Norway, project
numbers 324 183, 315 330, and 262 633. We acknowledge Elettra Sincrotrone
Trieste for providing access to its synchrotron radiation facilities and for
financial support under the IUS internal project (proposal 20215677). This
work has been partly performed in the framework of the nanoscience foundry and
fine analysis (NFFA-MUR Italy Progetti Internazionali) facility.

## Appendix: Methods

###  IV.1 Sample Preparation

A clean Bi(112) surface was prepared by repeated cycles of Ar+ ion sputtering
at $200$-$400\text{\,}\mathrm{eV}$, followed by annealing to
$T\approx$50\text{\,}\mathrm{\SIUnitSymbolCelsius}$$ for short duration. The
cleanliness of the surface was verified by sharp and oxide free Bi core levels
using X-ray photoelectron spectroscopy (XPS), and the crystallinty of the
surface by LEED.

###  IV.2 Band Structure and Spin Measurements

####  IV.2.1 Momentum Microscopy Measurements

Band structure measurements were performed at
$T\approx$115\text{\,}\mathrm{K}$$ using a NanoESCA III aberration corrected
EF-PEEM equipped with a focused helium discharge lamp primarily generating He
I photons at $h\nu=21.2$ eV, using pass energy
$E_{\text{P}}=$25\text{\,}\mathrm{eV}$$ and a $1.0\text{\,}\mathrm{mm}$
entrance slit to the energy filter. With the given settings, the instrument
had nominal energy and momentum resolutions of approximately
$100\text{\,}\mathrm{meV}$ and $0.02~{}\text{Å}^{-1}$, respectively.

Two-dimensional spin-polarised measurements at constant energies were
performed using an Ir spin filter coated with a monolayer of Au [28, 29]. The
spin filter measured the projection of the spin along
$\pm\mathbf{k}_{\text{y}}$, in-plane and along the 1D-line. The spin-
polarisation $P$ was calculated from the measurements using

$P=\frac{I_{\uparrow}-I_{\downarrow}}{S(I_{\uparrow}+I_{\downarrow})}$ (1)

where $I_{\uparrow}$ and $I_{\downarrow}$ are the intensities of the energy
surface when filtering spin along $+\mathbf{k}_{\text{y}}$ and
$-\mathbf{k}_{\text{y}}$, respectively [30]. A Sherman function $S=0.6$ was
assumed, based on preliminary calibration measurements [30, 31].

####  IV.2.2 High-Resolution Measurements

Higher energy resolution band structure measurements were performed at the
APE-LE endstation at Elettra Synchrotron, Italy, using a VG SCIENTA DA30
analyzer whilst cooling the Bi(112) crystal to $T=$77\text{\,}\mathrm{K}$$.

Spin measurements were performed using two three-dimensional vectorial spin-
polarimeters operated in the very low energy electron diffraction (VLEED)
regime [32]. From the spin signals detected by the two VLEED spin detectors it
is possible to reconstruct the full three-dimensional spin-vector carried by
the emitted photoelectrons. A Sherman function of $S=0.3$ was found from
calibration and was used in the analysis.

###  IV.3 Calculation Details

####  IV.3.1 Tight-Binding Calculations

The bulk band structure of bismuth was calculated within a TB model in Ref.
33. In this paper, we use the same model and parameters to calculate the band
structure of a semi-infinite system with a finite number of layers in the
$\hat{\mathbf{z}}$-direction and periodic boundary conditions in the
$\mathbf{\hat{x}\hat{y}}$-plane, where the sample has been oriented so that
the $\mathbf{\hat{x}\hat{y}}$-plane corresponds to the (112)-surface of
bismuth. By performing the Fourier transform for the in-plane coordinates, the
energy band structure can be calculated as eigenvalues of an effective 1D
problem at every in-plane momentum $\mathbf{k}$. From the eigenvector
corresponding to each eigenvalue at a given in-plane momentum, one may extract
various properties of the state, such as orbital content, spin, and the
spatial distribution in the direction perpendicular to the surface plane. The
latter can be used to discriminate between between bulk and surface states,
and is the basis for the choice of color in Fig. 2. Within the geometry in the
TB calculation, there are two surfaces, which both host surface states. In
Fig. 3, however, we only show the surface states corresponding to one of the
two surfaces, while the opposing surface contains surface states with opposite
spin-polarisation.

####  IV.3.2 DFT Calculations

First principles DFT calculations were carried out to better understand the
electronic structure and to further confirm the validity of the TB model. All
calculations were carried out with the QuantumESPRESSO DFT package using fully
relativistic pseudopotentials and projector augmented wave exchange-
correlation functionals (KJPAW). Bulk calculations were performed with a
$40\text{\,}\mathrm{Ry}$ plane wave cutoff and convergence threshold
$1\text{\times}{10}^{-8}\text{\,}\mathrm{Ry}$. The $\mathbf{k}$-points were
sampled using a Monkhorst-Pack grid of $12\times 12\times 12$ and the lattice
constant was found by a relaxation process. Surface states were calculated
using a slab geometry with 24 atomic layers and a separation of
$15\text{\,}\mathrm{\text{Å}}$ between slabs. Sampling of $\mathbf{k}$-points
was done using a Monkhorst-Pack grid of $10\times 10\times 1$. The cut-off
energy was $40\text{\,}\mathrm{Ry}$ and the convergence threshold
$1\text{\times}{10}^{-6}\text{\,}\mathrm{Ry}$.

## References

  * Hirohata  _et al._ [2020] A. Hirohata, K. Yamada, Y. Nakatani, I.-L. Prejbeanu, B. Diény, P. Pirro, and B. Hillebrands, Review on spintronics: Principles and device applications, J. Magn. Magn. Mater. 509, 166711 (2020).
  * Manipatruni  _et al._ [2019] S. Manipatruni, D. E. Nikonov, C.-C. Lin, T. A. Gosavi, H. Liu, B. Prasad, Y.-L. Huang, E. Bonturim, R. Ramesh, and I. A. Young, Scalable energy-efficient magnetoelectric spin-orbit logic, Nature 565, 35 (2019).
  * Brüne  _et al._ [2012] C. Brüne, A. Roth, H. Buhmann, E. M. Hankiewicz, L. W. Molenkamp, J. Maciejko, X.-L. Qi, and S.-C. Zhang, Spin polarization of the quantum spin Hall edge states, Nat. Phys. 8, 485 (2012).
  * König  _et al._ [2007] M. König, S. Wiedmann, C. Brüne, A. Roth, H. Buhmann, L. W. Molenkamp, X.-L. Qi, and S.-C. Zhang, Quantum Spin Hall Insulator State in HgTe Quantum Wells, Science 318, 766 (2007).
  * Roushan  _et al._ [2009] P. Roushan, J. Seo, C. V. Parker, Y. S. Hor, D. Hsieh, D. Qian, A. Richardella, M. Z. Hasan, R. J. Cava, and A. Yazdani, Topological surface states protected from backscattering by chiral spin texture, Nature 460, 1106 (2009).
  * Koo  _et al._ [2009] H. C. Koo, J. H. Kwon, J. Eom, J. Chang, S. H. Han, and M. Johnson, Control of Spin Precession in a Spin-Injected Field Effect Transistor, Science (New York, N.Y.) 325, 1515 (2009).
  * Bawden  _et al._ [2015] L. Bawden, J. M. Riley, C. H. Kim, R. Sankar, E. J. Monkman, D. E. Shai, H. I. Wei, E. B. Lochocki, J. W. Wells, W. Meevasana, T. K. Kim, M. Hoesch, Y. Ohtsubo, P. Le Fèvre, C. J. Fennie, K. M. Shen, F. Chou, and P. D. C. King, Hierarchical spin-orbital polarization of a giant Rashba system, Sci. Adv. 1, e1500495 (2015).
  * Hofmann [2006] P. Hofmann, The surfaces of bismuth: Structural and electronic properties, Prog. Surf. Sci. 81, 191 (2006).
  * Noguchi  _et al._ [2019] R. Noguchi, T. Takahashi, K. Kuroda, M. Ochi, T. Shirasawa, M. Sakano, C. Bareille, M. Nakayama, M. D. Watson, K. Yaji, A. Harasawa, H. Iwasawa, P. Dudin, T. K. Kim, M. Hoesch, V. Kandyba, A. Giampietri, A. Barinov, S. Shin, R. Arita, T. Sasagawa, and T. Kondo, A weak topological insulator state in quasi-one-dimensional bismuth iodide, Nature 566, 518 (2019).
  * Autès  _et al._ [2016] G. Autès, A. Isaeva, L. Moreschini, J. C. Johannsen, A. Pisoni, R. Mori, W. Zhang, T. G. Filatova, A. N. Kuznetsov, L. Forró, W. Van den Broek, Y. Kim, K. S. Kim, A. Lanzara, J. D. Denlinger, E. Rotenberg, A. Bostwick, M. Grioni, and O. V. Yazyev, A novel quasi-one-dimensional topological insulator in bismuth iodide $\beta$-$\text{Bi}_{4}\text{I}_{4}$, Nat. Mater. 15, 154 (2016).
  * Bianchi  _et al._ [2015] M. Bianchi, F. Song, S. Cooil, Å. F. Monsen, E. Wahlström, J. A. Miwa, E. D. L. Rienks, D. A. Evans, A. Strozecka, J. I. Pascual, M. Leandersson, T. Balasubramanian, P. Hofmann, and J. W. Wells, One-dimensional spin texture of Bi(441): Quantum spin Hall properties without a topological insulator, Phys. Rev. B 91, 165307 (2015).
  * Wells  _et al._ [2009] J. W. Wells, J. H. Dil, F. Meier, J. Lobo-Checa, V. N. Petrov, J. Osterwalder, M. M. Ugeda, I. Fernandez-Torrente, J. I. Pascual, E. D. L. Rienks, M. F. Jensen, and P. Hofmann, Nondegenerate Metallic States on Bi(114): A One-Dimensional Topological Metal, Phys. Rev. Lett. 102, 096802 (2009).
  * Hsieh  _et al._ [2009] D. Hsieh, Y. Xia, L. Wray, D. Qian, A. Pal, J. H. Dil, J. Osterwalder, F. Meier, G. Bihlmayer, C. L. Kane, Y. S. Hor, R. J. Cava, and M. Z. Hasan, Observation of Unconventional Quantum Spin Textures in Topological Insulators, Science (New York, N.Y.) 323, 919 (2009).
  * Xia  _et al._ [2009] Y. Xia, D. Qian, D. Hsieh, L. Wray, A. Pal, H. Lin, A. Bansil, D. Grauer, Y. S. Hor, R. J. Cava, and M. Z. Hasan, Observation of a large-gap topological-insulator class with a single Dirac cone on the surface, Nat. Phys. 5, 398 (2009).
  * Chen  _et al._ [2009] Y. L. Chen, J. G. Analytis, J.-H. Chu, Z. K. Liu, S.-K. Mo, X. L. Qi, H. J. Zhang, D. H. Lu, X. Dai, Z. Fang, S. C. Zhang, I. R. Fisher, Z. Hussain, and Z.-X. Shen, Experimental Realization of a Three-Dimensional Topological Insulator, $\text{Bi}_{2}\text{Te}_{3}$, Science 325, 178 (2009).
  * Pan  _et al._ [2011] Z.-H. Pan, E. Vescovo, A. V. Fedorov, D. Gardner, Y. S. Lee, S. Chu, G. D. Gu, and T. Valla, Electronic Structure of the Topological Insulator $\mathrm{Bi}_{2}\mathrm{Se}_{3}$ Using Angle-Resolved Photoemission Spectroscopy: Evidence for a Nearly Full Surface Spin Polarization, Phys. Rev. Lett. 106, 257004 (2011).
  * Wells  _et al._ [2008] J. W. Wells, K. Handrup, J. F. Kallehauge, L. Gammelgaard, P. Bøggild, M. B. Balslev, J. E. Hansen, P. R. E. Petersen, and P. Hofmann, The conductivity of Bi(111) investigated with nanoscale four point probes, J. Appl. Phys. 104, 053717 (2008).
  * Guo  _et al._ [2011] H. Guo, K. Sugawara, A. Takayama, S. Souma, T. Sato, N. Satoh, A. Ohnishi, M. Kitaura, M. Sasaki, Q.-K. Xue, and T. Takahashi, Evolution of surface states in $\mathrm{Bi}_{1-x}\mathrm{Sb}_{x}$ alloys across the topological phase transition, Phys. Rev. B 83, 201104 (2011).
  * Zhu and Hofmann [2014] X.-G. Zhu and P. Hofmann, Topological surface states on $\text{Bi}_{1-x}\text{Sb}_{x}$: Dependence on surface orientation, termination, and stability, Phys. Rev. B 89, 125402 (2014).
  * [20] See Supplementary Material at [insert link] which includes Ref. [34] for details of the DFT calculations, how the binding energy of the crossing points of the surface states were found, the orbital nature of the electronic states and details regarding the measured out-of-plane spin-component and the spin-vector.
  * Damascelli [2004] A. Damascelli, Probing the Electronic Structure of Complex Systems by ARPES, Phys. Scr. 2004, 61 (2004).
  * Kevan and Eberhardt [1992] S. D. Kevan and W. Eberhardt,  _Angle-Resolved Photoemission: Theory and Current Applications_ , edited by S. D. Kevan (Elsevier Amsterdam, 1992) Chap. 4.
  * Day  _et al._ [2019] R. P. Day, B. Zwartsenberg, I. S. Elfimov, and A. Damascelli, Computational framework chinook for angle-resolved photoemission spectroscopy, npj Quantum Materials 4, 10.1038/s41535-019-0194-8 (2019).
  * Bakkelund [2022] J. Bakkelund,  _Investigating spin-split surface states on semimetals and correlation effects in a ferromagnet_ , Master’s thesis, Norwegian University of Science and Technology (NTNU), https://hdl.handle.net/11250/3007070 (2022).
  * Crain  _et al._ [2004] J. N. Crain, J. L. McChesney, F. Zheng, M. C. Gallagher, P. C. Snijders, M. Bissen, C. Gundelach, S. C. Erwin, and F. J. Himpsel, Chains of gold atoms with tailored electronic states, Phys. Rev. B 69, 125401 (2004).
  * Kim  _et al._ [2005] T. K. Kim, J. Wells, C. Kirkegaard, Z. Li, S. V. Hoffmann, J. E. Gayone, I. Fernandez-Torrente, P. Häberle, J. I. Pascual, K. T. Moore, A. J. Schwartz, H. He, J. C. H. Spence, K. H. Downing, S. Lazar, F. D. Tichelaar, S. V. Borisenko, M. Knupfer, and P. Hofmann, Evidence against a charge density wave on Bi(111), Phys. Rev. B 72, 085440 (2005).
  * Leuenberger  _et al._ [2013] D. Leuenberger, H. Yanagisawa, S. Roth, J. H. Dil, J. W. Wells, P. Hofmann, J. Osterwalder, and M. Hengsberger, Excitation of Coherent Phonons in the One-Dimensional Bi(114) Surface, Phys. Rev. Lett. 110, 136806 (2013).
  * Tusche  _et al._ [2015] C. Tusche, A. Krasyuk, and J. Kirschner, Spin resolved bandstructure imaging with a high resolution momentum microscope, Ultramicroscopy Special Issue: LEEM-PEEM 9, 159, 520 (2015).
  * Kirschner  _et al._ [2013] J. Kirschner, F. Giebels, H. Gollisch, and R. Feder, Spin-polarized electron scattering from pseudomorphic Au on Ir(001), Phys. Rev. B 88, 125419 (2013).
  * Meier  _et al._ [2009] F. Meier, J. H. Dil, and J. Osterwalder, Measuring spin polarization vectors in angle-resolved photoemission spectroscopy, New J. Phys. 11, 125008 (2009).
  * Sherman [1956] N. Sherman, Coulomb Scattering of Relativistic Electrons by Point Nuclei, Phys. Rev. 103, 1601 (1956).
  * Bigi  _et al._ [2017] C. Bigi, P. K. Das, D. Benedetti, F. Salvador, D. Krizmancic, R. Sergo, A. Martin, G. Panaccione, G. Rossi, J. Fujii, and I. Vobornik, Very efficient spin polarization analysis (VESPA): new exchange scattering-based setup for spin-resolved ARPES at APE-NFFA beamline at Elettra, J. Synchrotron Radiat. 24, 750 (2017).
  * Liu and Allen [1995] Y. Liu and R. E. Allen, Electronic structure of the semimetals Bi and Sb, Phys. Rev. B 52, 1566 (1995).
  * Himpsel [1980] F. J. Himpsel, Experimental determination of bulk energy band dispersions, Appl. Opt. 19, 3964 (1980).


# MemServe: Context Caching for Disaggregated LLM Serving with Elastic Memory
Pool

Cunchen Hu2,3, Heyang Huang2,3, Junhao Hu4, Jiang Xu1, Xusheng Chen1, Tao
Xie4,
Chenxi Wang2,3, Sa Wang2,3, Yungang Bao2,3, Ninghui Sun2,3, Yizhou Shan1
1Huawei Cloud, 2UCAS 3ICT, CAS 4Peking University

###### Abstract.

Large language model (LLM) serving has transformed from stateless to stateful
systems, utilizing techniques like context caching and disaggregated
inference. These optimizations extend the lifespan and domain of the KV cache,
necessitating a new architectural approach. We present MemServe, a unified
system that integrates both inter-request and intra-request optimizations.
MemServe introduces MemPool, an elastic memory pool managing distributed
memory and KV caches across serving instances. Using MemPool APIs, MemServe
combines context caching with disaggregated inference for the first time,
supported by a global scheduler that enhances cache reuse through a global
prompt tree-based locality-aware policy. Tests show that MemServe
significantly improves job completion time and time-to-first-token.

##  1\. Introduction

Large language models (LLMs) and their underlying transformer architecture
have revolutionized AI, becoming foundational to many emerging applications
and a crucial workload in data centers. While high-quality models are
essential, it is equally important to serve these models on a massive scale at
a reasonably low cost. As a result, numerous approaches have been proposed to
enhance the cost-efficiency of LLM serving, such as context caching (sglang, ;
pensieve, ), disaggregated inference (zhong2024distserve, ;
patel2023splitwise, ), and sequence parallelism (infiniteLLM, ).

As a result, LLM serving has evolved from a stateless to a stateful system,
leveraging dependencies inherent in inference requests. These dependency-
exploiting techniques can be classified into two types: inter-request and
intra-request. Inter-request techniques exploit dependencies across requests.
The notable one is context caching (sglang, ), which reuses the KV cache for
requests that share the same prompt prefix, thereby speeding up the prefill
phase. Intra-request techniques, on the other hand, exploit dependencies
within a single request. Two prominent examples are disaggregated inference,
which splits a request into two sub-requests for better scheduling
(patel2023splitwise, ), and sequence parallelism, which divides a request into
multiple sub-requests to distribute load (infiniteLLM, ).

A common theme in these dependency-exploiting techniques is that they require
novel logic to manage and transfer the KV cache, which is the intermediate
data produced during LLM inference. Inter-request methods preserve the KV
cache across requests, extending its lifetime from a single request to
potentially infinite. Intra-request methods transfer the KV cache across
multiple inference instances, extending its domain from a single instance to
distributed instances. However, deploying a stateful LLM serving system with
these optimizations is challenging due to conflicting or missing mechanisms
for managing the LLM’s intermediate KV cache data. We have identified two key
problems.

The first problem is that LLM serving systems cannot simultaneously apply any
existing inter-request and intra-request dependency-exploiting optimizations.
Current context caching (inter-request) methods are designed without
considering intra-request scenarios. As a result, disaggregated inference
(intra-request) cannot benefit from context caching because it lacks the
mechanisms to utilize the KV cache from decode back to prefill instances for
future reuse. Similarly, sequence parallelism distributes the KV cache across
multiple instances and lacks the mechanisms and algorithms needed to preserve
and reuse it. This issue arises because intra-request techniques break a
tightly coupled request into multiple loosely coupled sub-requests,
complicating KV cache management in a distributed setting.

The second problem is that LLM serving systems lack a holistic, top-down
design to effectively utilize existing inter-request techniques. Context
caching benefits from reusing historical KV cache by running requests that
share a common prefix in the same serving instance. However, current LLM
serving systems schedule requests across multiple serving instances based on
load or session IDs, which fails to maximize KV cache reuse across sessions.

These issues arise because existing LLM serving systems are built on the
assumption that the KV cache is merely intermediate data scoped to a single
request on a single instance. With emerging dependency-exploiting techniques,
the lifespan of the KV cache has been extended, and its management has
expanded to a distributed setup. This paradigm shift calls for a fundamental
rethinking of LLM serving architectures.

In this work, we propose Memory-enhanced model Serving, or MemServe, to handle
inter-request and intra-request optimizations within a unified system. To
tackle the challenges of managing the KV cache across distributed instances,
MemServe introduces an elastic memory pool, or MemPool, which is a substrate
for managing all cluster memory, including CPU DRAM and GPU HBM. MemPool
offers a rich set of APIs for managing distributed memory and KV cache.
Utilizing these APIs, MemServe implements context caching over standard
prefill-decode-colocated (PD-colocated) instances (sglang, ) and disaggregated
inference (zhong2024distserve, ; tetriserve-2024, ). Moreover, MemServe
enhances disaggregated inference with context caching, reaping both benefits.
Finally, to maximize KV cache reuse, MemServe employs a global scheduler that
incorporates a locality-aware policy using novel global prompt trees for best-
effort routing.

The MemPool is a core component of MemServe, providing three types of APIs:
memory, indexing, and distributed data transfer. It runs within each inference
instance, managing all local memory with a fixed-size memory allocator. The
indexing APIs are crucial for building context caching. MemPool uses an
internal index to map prompt tokens to the KV cache, managing both the active
KV cache for ongoing requests and the historical KV cache retained after
requests are completed. The MemPool offers a simple data transfer API that
abstracts three heterogeneities: parallelism, network, and memory medium. As a
unified platform, MemPool supports all known inter-request and intra-request
optimizations as well as any combinations (see Figure 3).

MemServe bridges the gap between context caching (inter-request) and
disaggregated inference (intra-request) in four steps using MemPool APIs: (a)
we first use a distributed API to reproduce disaggregated inference, (b) we
then add caching to prefill-only instances using index APIs, (c) we apply the
same caching to decode-only instances, (d) finally we enable decode-to-prefill
data transfer, as illustrated in Figure 4. However, it is challenging to hit
two birds with one stone. We observed increasing overheads due to naive
discrete memory layouts and point-to-point network primitives from existing AI
network stacks. To address this, MemServe proposes co-optimizing memory layout
and network transfer using huge pages.

We implement MemPool and global scheduler from scratch, 5.6K SLOC in Python
and 1.6K SLOC in C++ We modify vLLM (vllm-sosp23, ) to build context caching
with disaggregated inference, 200 SLOC in Python and 400 SLOC in CUDA C++ We
use NCCL send and recv pairs for data transmission between GPUs and socket if
any side is DRAM.

We run all tests atop a single server with eight H800-80G GPUs. We evaluate
MemServe across four settings: (1) PD-colocated, (2) PD-colocated with
caching, (3) PD-disaggregated, and (4) PD-disaggregated with caching. The
first setting runs a vanilla vLLM. The last three settings are MemServe
running adapted vLLM using MemPool APIs. While running ShareGPT workload
(sharegpt, ), the PD-disaggregated with caching setting outperforms others.
Specifically, MemPool-based disaggregated inference improves JCT by up to 42%
compared to PD-colocated. Enhancing disaggregated inference with context
caching can further improve JCT by 29%! When executing the LooGLE dataset,
which features extended prompts and relatively short generation lengths,
disaggregated inference boosts JCT by up to 10.8% compared to PD-colocated
setups. Additionally, context caching offers further enhancements, potentially
improving JCT by 26.9%.

In summary, we make the following contributions:

  * •

We propose MemPool, a memory pool designed for LLM serving with a rich set of
APIs.

  * •

We build the first disaggregated inference with context caching in MemServe
based on MemPool APIs.

  * •

We propose a novel prompt trees-based locality-aware policy for scheduling LLM
requests.

##  2\. Background

Generative LLM Inference LLM inference involves generating a sequence of
output tokens in response to an input prompt. This process consists of two
distinct phases: prefill and decode. During the prefill phase, the model
processes the prompt to generate the key-value (KV) cache. The KV cache
comprises key-value pairs generated in the self-attention mechanism. In the
decode phase, the model uses the KV cache to generate tokens iteratively. The
size of the KV cache grows linearly with increasing number of generated
tokens.

Inter-Request Optimization This type of optimization exploits dependencies
among requests for better performance. Context caching is the only known
technique in this category. To build context caching, the model stores and
reuses the KV cache from the self-attention mechanism to avoid redundant
computations across similar or repeated requests. This is useful in scenarios
where multiple requests share common prefixes or contexts. Two mechanisms are
essential. First, an index is required to find dependencies among requests and
consequently find the preserved KV cache (see Table 2). Second, a modified
inference engine and attention kernel to reuse the historical KV cache (see
SGLang (sglang, )).

Intra-Request Optimization This type of optimization exploits dependencies
within a request to enhance performance. Two notable examples are
disaggregated inference (patel2023splitwise, ; zhong2024distserve, ;
tetriserve-2024, ; strati2024dejavu, ) and sequence parallelism (infiniteLLM,
). Generally, disaggregating prefill from decode reduces interference between
these two stages and allows each to scale independently with heterogeneous
hardware. However, this breaks a single request into two sub-requests and
requires rigorous KV cache transmission from prefill to decode. The same goes
for sequence parallelism, in which distributed instances need to exchange the
outputs of self-attention in a rigorous manner. Overall, intra-request
optimization demands efficient mechanisms for transferring the KV cache among
instances.

##  3\. MemServe Overview

Figure 1. MemServe Architecture. It supports three types of inference
instances: prefill-only, decode-only, and PD-colocated. Each inference engine
runs over one or multiple AI servers, depending on the parallelism
configuration.

MemServe is designed as a large-scale LLM serving system that efficiently
handles inter-request and intra-request optimizations. It comprises three main
components: a global scheduler, multiple types of inference instances, and an
elastic memory pool (MemPool), as shown in Figure 1. The MemPool offers a set
of APIs for memory allocation, index management, and distributed transfer
(§4). MemServe builds context caching atop both regular and disaggregated
inference architectures using MemPool APIs (§5). The global scheduler forwards
inference requests from users to the right inference instance. It uses
locality-aware policies based on novel distributed prompt trees, maximizing KV
cache reuse (§6).

##  4\. Elastic Memory Pool

The MemPool manages all memory in the inference cluster, including CPU DRAM
and GPU HBM. MemPool runs within each inference instance, collectively
offering a set of distributed memory pool APIs (§4.1). It manages both the
active KV cache used by ongoing requests and the historical KV cache retained
after requests are completed. An indexing layer maps prompt tokens to the
historical KV cache (§4.2), ensuring efficient retrieval of cached data. The
MemPool has efficient mechanisms for exchanging data between instances,
alleviating inference engines from dealing with heterogeneous hardware (§4.3).
Overall, this design makes MemPool a versatile and generic platform capable of
supporting both intra-request and inter-request optimizations within a unified
system (§4.5).

###  4.1. API

Table 1. Elastic Memory Pool APIs. Type can be HBM-only, DRAM-only, or mixed. Each address encodes instance ID. Transfer flags can control on-demand allocation. API |  Parameters |  Description
---|---|---
alloc_mem |  size, type, id |  alloc a certain type of memory on a given instance (@id), return addrList
free_mem |  addrList |  free memory
insert |  tokenList, addrList, flags |  insert prompt token and KV cache address mapping into local index
match |  tokenList |  find prompt’s cached data if any, return addrList
delete |  tokenList |  delete prompt’s cached data if any
swap_out |  num_blocks |  swap a given number of blocks from HBM to DRAM
swap_in |  addrList |  swap blocks with given address from DRAM to HBM
transfer |  id, srcAddrList, dstAddrList, flags, private |  transfer data to the specified instance (@id), dstAddrList is optional, flags control behaviors at the destination, and private carries user data
transfer_ with_insert |  id, tokenList, srcAddrList, dstAddrList, flags, private |  transfer tokenList and its cached data to the specified instance. The receiver will call an extra insert.

We show MemPool APIs in Table 1, broadly divided into three categories: memory
block, index, and distributed transfer. The inference engine can use memory
block APIs to allocate fixed-size memory blocks for storing KV cache or other
data. The engine can also call the index APIs for context caching solutions.
For example, once requests are finished, the engine can call insert to
transition the active KV cache into the historical KV cache and create a
mapping from prompt tokens to the KV cache. The engine can invoke distributed
APIs, such as transfer, to exchange the KV cache across instances when
building inter-request optimizations. Overall, the MemPool provides a rich
framework for managing distributed memory and implementing efficient context
caching and data exchange mechanisms.

###  4.2. Indexing

Table 2. Compare Indexing Methods. MemServe’s MemPool uses prompt tokens for its generality. Indexing |  Description
---|---
Token ID |  Use prompt tokens. Universally applicable.
Session ID |  Use client-server session ID. Limited scope.
Document ID |  Use document file ID. Limited scope.

The MemPool has an index layer to map prompt tokens to the historical KV
cache. MemPool traverses the index whenever engines call insert, match,
delete, etc. The LLM serving world has three indexing methods: token, session,
and document IDs (see Table 2). Token-based indexing is known for its
generality, as it works for any shared prompt-prefix cases (sglang, ). The
session and document ID indexing are simpler but can only reuse shared prompts
within a chat session or across sessions using the same document
(wu2024loongserve, ; gemini-context-caching, ). We adopt the token-based
indexing method for broad applicability. To implement this index, MemPool
utilizes the radix tree proposed by SGLang (sglang, ), with two key
extensions. First, because MemPool manages both GPU HBM and CPU DRAM, we
enable the radix tree to reference data located anywhere in the system.
Second, since we also use the radix tree to build the global prompt tree in
the global scheduler (§6), we add a field to indicate which inference instance
holds the data. Note that while mixed indexing methods are possible, we will
explore this in future work.

To minimize data reshaping overhead, we maintain the original memory layout
when transitioning the active KV cache to the historical KV cache before
inserting it into MemPool. Consequently, MemPool’s indexing granularity aligns
with the inference engines’ configuration. For example, in our tests with
vLLM, which uses a block size of 16 tokens, our radix tree nodes point to KV
cache blocks of 16 tokens.

###  4.3. Distributed Transfer

Figure 2. MemPool Transfer API. The left shows the workflow of transfer and
transfer_with_insert. The right shows asymmetric parallelism and memory
medium.

The MemPool exposes distributed APIs for exchanging data among inference
instances. They serve as the building blocks for intra-request or inter-
request dependency-exploiting techniques. Our design rationale is to expose
simple APIs that mask the underlying heterogeneity from inference engines.

Figure 2 shows the transfer workflow and our approach to handling
heterogeneity. We break down the workflow into three steps: allocation,
transmission, and insertion. When the sender inference instance initiates a
transfer, it sends a request to the receiver inference instance. Upon
receiving this request, the receiver invokes alloc_mem locally to allocate HBM
or DRAM based on the type specified by the sender. The receiver then returns
the allocated address list and its parallelism configuration to the sender,
completing the allocation step. Then, the sender transmits the KV cache to the
receiver using the fastest available path. Once all data is received, the
receiver notifies the sender, completing the transmission step. Next, The
receiver checks whether this is a transfer_with_insert call. If so, it invokes
the insert function locally to insert the newly transmitted prompt tokens and
historical KV cache into its local index, completing the insertion step.
Finally, the sender completes the transfer API call once the receiver returns
ok.

We propose the transfer_with_insert as it can avoid an extra network round-
trip for establishing mapping, which is particularly useful for transferring
historical KV cache from a decode-only instance to a prefill-only instance.

Additionally, users can call the transfer API with a specific destination
address list, allowing them to skip the initial allocation step. This feature
is particularly useful for constructing layer-by-layer transmissions in
disaggregated inference (see Figure 5).

The transmit step is the most challenging as it must deal with three types of
heterogeneities between the sender and the receiver: parallelism, memory, and
network. To manage asymmetric parallelism, the sender first checks how the KV
cache is partitioned along tensor-parallel or pipeline-parallel dimensions.
Once determined, the sender partitions its local cache and invokes the
appropriate network primitives (top-right in Figure 2). Memory asymmetry can
occur if the historical KV cache has been swapped out to DRAM (bottom-right in
Figure 2). MemPool always tries to transmit data using the fastest link with
the least data copies. But this is highly hardware-dependent. If MemPool uses
the latest hardware, such as NVIDIA SuperPOD, where all HBM and DRAM are
connected by high-speed NVLINK, handling memory asymmetry is as simple as
performing a memory copy. However, on regular GPU servers, additional memory
copies in the data path are inevitable. While implementing MemPool distributed
APIs, we find existing network primitives ill-fit for handling emerging LLM
inference workloads. We will discuss this in §7.

###  4.4. Failure Handling and Scaling

We discuss how MemPool handles failures and dynamic scaling during runtime. As
Figure 1 shows, MemPool is deployed as part of an inference instance. Hence,
the failure and scaling granularity is a single instance, which can be one or
multiple AI servers, depending on the parallelism configuration.

MemServe has a cluster management (CM) module as shown in Figure 1, which is a
centralized service for maintaining cluster configuration. The CM is
responsible for adding or removing instances and monitoring cluster health. In
our current design, memory block and distributed transfer APIs can modify the
states of remote instances. When an instance fails, any in-flight requests
from other instances will time out. The CM detects such failures through
regular heartbeats and broadcasts updated cluster information to all running
instances. Upon receiving this notification, each instance releases any memory
blocks allocated by the failed instance to prevent memory leaks.

###  4.5. Use Cases

Figure 3. Use Cases Enabled By MemPool. Circle 1 is context caching. Circle 2 is disaggregated inference. Circle 3 is sequence parallelism. Solid gray lines mean MemPool index API calls. Solid red lines mean MemPool distributed APIs. MemPool enables all use cases in one platform. Table 3. Atomic Serving Scenarios Supported by MemPool. As a unified platform, MemPool supports any combo of inter-request or intra-request optimizations in one system. Scenario | Type |  APIs Used
---|---|---
Context Caching | inter |  index (insert,match,delete,evict,etc)
Disagg. Inference | intra |  dist (transfer,transfer_with_insert)
Sequence Parallel | intra |  dist (transfer)
Request Migration | N/A |  dist (transfer)

MemPool is a versatile and generic platform designed to support both inter-
request and intra-request dependency-exploiting techniques. Figure 3
illustrates how various existing inter-request and intra-request optimizations
can be integrated into a unified system using MemPool APIs. Table 3 lists the
advanced APIs employed in these optimizations. To build context caching atop
regular PD-colocated inference instances, one can use index APIs such as
insert and match. To build disaggregated inference, one can call transfer API
to send active KV cache from a prefill-only to a decode-only instance. To
build sequence-parallel (SP) inference, one can use the transfer API to
exchange attention outputs among instances, akin to InfiniteLLM (infiniteLLM,
).

What sets MemPool apart is its ability to seamlessly enable these
optimizations within a single system using a common set of APIs. Next, we will
showcase how to enhance disaggregated inference with context caching. We leave
other combinations for future work.

##  5\. Caching for Disaggregated Inference

Figure 4. Enhancing Disaggregated Inference with Context Caching using MemPool APIs. The engine box means an adapted inference engine such as vLLM. Circled numbers mean steps taken to build the solution. A-KV is active KV cache. H-KV is historical KV cache. Table 4. Towards Full-Fledged Context Caching in Disaggregated Inference. Refer to Figure 4 for step numbers. Design | Steps | Description
---|---|---
PD-Basic | 1 | Basic PD, no caching
PD-Caching-1 | 1+2 | Caching at P
PD-Caching-2 | 1+2+3+4 | Caching at D
PD-Caching-3 | 1+2+3+4+5 | Full-fledged caching

Context caching exploits dependency across requests, while disaggregated
inference exploits dependency within a request. However, they fail to coexist
due to missing mechanisms around KV cache management. We enhance disaggregated
inference with context caching using MemPool APIs. To the best of our
knowledge, this is the work introducing caching to disaggregated inference.

###  5.1. Design

We show how to gradually build towards a full-fledged design in four design
milestones in Table 4, utilizing five key MemPool APIs as highlighted in
Figure 4.

(a) PD-Basic. This is the basic disaggregated inference architecture proposed
by DistServe (zhong2024distserve, ) and Spliwise (patel2023splitwise, ). To
realize this design, we make minor changes to an existing inference engine
(e.g., vLLM (vllm-sosp23, )). As a result, the prefill instance will call
MemPool’s transfer API to transfer the active KV cache produced after the
prefill phase to the decode instance. We carry essential metadata the decode
instance requires in transfer’s private field, such as request ID, sampling
parameters, prompt tokens, etc.

(b) PD-Caching-1. This is the first caching-enhanced disaggregated inference
design. We enable caching at the prefill-only instance by calling insert to
retire the active KV cache as the historical KV cache such that future
inferences can utilize the saved data to reduce recomputation (step 2 in
Figure 4). This caching design only preserves historical KV cache produced by
the prefill phase but none from the decode phase, so it works well for
workloads that share long common prefix prompts, e.g., system prefix (sglang,
). The major drawback of this design is that in a multi-turn chat scenario
(e.g., document QA (li2023loogle, )), the prefill-only instance needs to
repeatedly forward the same set of active KV cache to the decode-only
instance, wasting bandwidth and affecting the time-to-second-token. We
therefore propose the next design milestone to address this issue.

(c) PD-Caching-2. This design enables caching at the decode-only instance to
reduce repeated data movement. We make two key changes atop PD-Caching-1.
First, the prefill-only instance now calls transfer_with_insert instead of
transfer such that the decode-only instance will insert the transmitted KV
cache produced by the prefill phase into its local index (§4.3). Second, after
a request finishes, the decode-only instance calls insert to preserve the KV
cache produced by the decode phase into its local index. With the help of
locality-aware scheduling (§6), the prefill-only instance now only needs to
transfer new KV cache data incrementally. Though this design reduces data
movement from prefill-only to decode-only instances, it does not improve
context caching at the prefill-only instance since it lacks the historical KV
cache from the decode phase. As a result, the benefit of context caching stays
flat with increasing prompt in a multi-turn chat scenario. We therefore
propose the next design milestone to address this issue.

(d) PD-Caching-3. This design enables full-fledged context caching for
disaggregated inference architecture. We make one change atop PD-Caching-2:
after a request finishes, the decode-only instance calls transfer_with_insert
to transmit the KV cache produced by the decode phase to the prefill-only
instance (step 5 in Figure 4). As a result, the prefill-only instance’s
preserved historical KV cache grows, and the benefit of context caching
increases linearly with the number of turns.

In all, we illustrate how MemPool’s simple APIs can be used to build a range
of advanced solutions, from basic disaggregated inference to full-fledged
context caching. Nevertheless, MemPool only provides primitves for
transferring data among inferences and managing local data. It is up to the
users (e.g., inference engines) to decide the memory layout and the number of
API calls. Next, we will discuss design challenges around memory and network.

###  5.2. Memory and Network Optimization

Figure 5. Optimize Network&Memory for Disagg. Inference.

We discuss how memory and network play a key role in building context caching
with disaggregated inference. As Splitwise (patel2023splitwise, ) points out,
there are two ways to transfer active KV cache from prefill to decode
instance: by-layer or by-request. The by-layer approach transfers the KV cache
once a layer has finished computation. The by-request approach transfers the
KV cache once the prefill phase is completed. Splitwise found that by-layer
outperforms by-request because it overlaps computation and communication,
hence speeding up time-to-second token (or TTST). We make the same observation
when the load is low. However, both incur non-trivial overhead with increasing
load due to excessive network transfers. We find the root causes are (1)
discrete memory layout and (2) inadequate network primitives.

Paging-based dynamic memory management introduced by PagedAttention (vllm-
sosp23, ) is now the de facto standard in LLM serving systems, e.g., in vLLM
(vllm-sosp23, ), TensorRT-LLM (tensorrt-llm, ). Regardless of where the paging
mechanism is implemented (engine (vllm-sosp23, ) or driver
(prabhu2024vattention, )), the KV cache is partitioned and stored in fixed-
sized memory blocks. The block size is configurable, usually in the number of
tokens, e.g., 8 tokens worth of KV cache per block. Existing engines manage
the KV cache in a fine-grained manner. For example, vLLM allocates two blocks
per LLM layer. Given an LLM model with $L$ layers and 8 tokens per block, the
engine needs $2*L$ blocks to store the KV cache of 8 tokens. Although paging
improves utilization (vllm-sosp23, ), the discrete memory layout presents huge
challenges when implementing disaggregated inference using existing AI network
stacks.

The de facto network stack in AI is collective libraries such as NCCL (url-
nccl, ). These libraries work best for typical AI workloads using tensor or
pipeline parallelism, but they fall short in supporting LLM serving’s intra-
request optimizations such as disaggregated inference (tetriserve-2024, ) or
sequence parallel (infiniteLLM, ). These new patterns require efficient point-
to-point, gather, and scatter primitives between HBM or DRAM, similar to RDMA
verbs (rdma-verbs, ). As discussed in §7, we implement transfer using NCCL
send and recv APIs, and each call only transmits a single block. Since the KV
cache is discrete, the number of network API calls equals the number of
discrete memory blocks, regardless of whether the by-layer or by-request
approach is used. This is the root cause of why both incur overhead with
increasing load.

To address challenges caused by paging and poor network primitives, we propose
to reduce fragmentation by aggregating smaller KV blocks into large ones, akin
to using huge pages. Specifically, instead of having two blocks per layer, we
aggregate them into one block; the new block size equals $2*L$ smaller blocks.
This effectively reduces the number of network API calls by $2*L$ times. This
optimization works only for the by-request approach, as the by-layer approach
inevitably needs to call the network APIs at least $L$ times. Our test shows
this technique improves network performance alone by a large margin, as shown
in Figure 11.

We compare by-layer, by-request, and by-request-agg (proposed optimization) in
Figure 5 across memory layout and transmission timeline. Our test shows that
under low load, by-layer achieves the lowest JCT, but under high load, by-
layer-agg outperforms by-layer thanks to reduced network calls, as shown in
Figure 12.

###  5.3. Cost Model for Context Caching

Table 5. Context Caching Cost Model Factors. Factor | Description
---|---
Prompt-Length | The length of the current prompt.
Cached-Ratio | The ratio of cached tokens
Cached-Locations | Historical KV cache locations.
Batch-Size | The running batch size.
Block-Size | Size of paging blocks.

We design a cost model for prefill-only instances (and PD-colocated instances)
to decide whether context caching is beneficial. The cost model concerns
factors listed in Table 5. Our cost model employs fitted curves derived from
experimental data in Figure 13.

Requests arrive at MemServe go through the following steps. First, requests
pass through the global scheduler’s locality-aware scheduler (described in
§6), which determines the instance(s) to run the request and locates the
historical KV cache. Next, requests arrive at a specific instance for prefill.
We call MemPool’s match API to check whether the new request has locally
cached data. If so, the match returns a list of local block addresses. Along
with the historical KV cache list sent from the global scheduler, we now know
the cache ratio and cache locations for a specific request. Then, instead of
invoking the cost model for each request individually, the instance’s local
request scheduler groups multiple requests to consider the Batch Size factor
for a single cost model invocation. This is because the same cache ratio has
varied benefits depending on the batch’s prompt length. Finally, assuming the
cost model decides to reuse context caching for a batch of requests, the
instance’s scheduler moves missing data from its DRAM or remote instances to
its HBM using MemPool APIs. The instance starts the prefill phase once all
historical KV cache is in the HBM.

##  6\. Locality-Aware Global Scheduling

Figure 6. Global Scheduler Architecture. We highlight the global prompt trees-based locality-aware scheduling, it consists of a lookup path (left) and an update path (right). Table 6. Global Request Scheduling Policies. We compare whether they improve inter-session or intra-session context caching. A session can be an HTTP session. Name | Intra-Session | Inter-Session
---|---|---
Least Load | N | N
Session-ID-Based | Y | N
Prompt-Tree-Based | Y | Y

In this section, we describe MemServe’s global scheduler (GS). The GS routes
requests from external services to underlying inference instances and returns
generated responses in a streaming fashion. To improve context caching at a
large scale, we propose global prompt trees and a locality-aware scheduling
policy. Figure 6 shows GS’s architecture.

Global Prompt Trees. Since MemServe runs three types of inferences, the GS
employs three types of prompt trees, for prefill-only, decode-only, and PD-
colocated instances. Each tree type has a set of radix trees, the same as the
ones used by MemPool with one extra field per tree node pointing to the
instance storing the KV cache. The global prompt trees support regular insert
and match APIs as listed in Table 1. For now, both GS’s global prompt tree and
each inference’s local prompt tree share the same indexing granularity.

Scheduling. When a request arrives at the GS, it goes through the following
steps. First, the GS runs a tokenizer to turn prompt strings into token IDs.
Second, the GS queries the global prompt tree by concurrently calling match
against all types of trees. Third, the GS sends query results along with
current load info to a policy module. The policy module chooses an instance
with the longest common prefix (i.e., the largest preserved historical KV
cache). Once the instance is chosen, the GS checks whether there exist
instances storing extra historical KV cache that is not present in the chosen
instance. If so, the policy engine also outputs a list of such instances and
the corresponding token IDs. Finally, the GS sends the request and metadata to
the chosen instance. We update the global prompt trees when instances return
responses back to callers.

Discussion. (1) Our proposed prompt-tree-based policy is a best-effort
scheduling policy. It tries to maximize context caching reusing opportunities.
Since the GS only updates its prompt tree when responses pass through it, the
GS is unaware of local eviction events in underly instances. Therefore, the
GS’s local prompt tree can be outdated. We address this issue by configuring
the global prompt trees with a time-to-live (TTL), commonly in minutes. (2) We
compare three global request scheduling policies in Table 6 across two
dimensions. The least-load policy means selecting the instance with the least
load, unaware of any locality. The session-ID-based policy routes requests
based on a connection ID (e.g., HTTP session). This policy enables context
caching within a session. Our prompt-tree-based policy can exploit caching
opportunities across sessions, reusing most context caching.

##  7\. MemServe Implementation

MemServe has three key parts: MemPool, context caching with disaggregated
inference, and global scheduler. We implement MemPool from scratch, 5K SLOC in
Python and 1.6K SLOC in C++. We modify vLLM (vllm-sosp23, ) to build context
caching with disaggregated inference, 200 SLOC in Python and 400 SLOC in CUDA
C++. The global scheduler and the cluster management have 600 SLOC in Python.

MemPool. It has two parts: a Python-based library that exposes API to the
inference engine and a C++ core part that executes data transmission.
Currently, MemPool uses NCCL’s send and recv point-to-point APIs to transmit
data between HBM and uses socket API if any side contains DRAM. We have not
implemented RDMA-based transmission because we only have a single AI machine.
As we’ve mentioned earlier, NCCL is a collective library designed for typical
tensor and pipeline parallel AI workloads but not ideal for point-to-point
communication. Specifically, its send and recv APIs only specify source
addresses but no destination addresses. Hence, ensuring ordering between a
sender and a receiver is challenging, especially if we aim to achieve high
parallelism using multiple threads. As a result, we end up using a single
thread per NCCL communicator to ensure ordering. Additionally, as NCCL has no
gather or scatter APIs, we call the send-recv API pair multiple times to
transmit data across heterogeneous parallelism instances (Figure 2).

Context Caching with Disaggregated Inference. We adapt vLLM (vllm-sosp23, ) to
using MemPool APIs. Specifically, we replace its original cache engine and
hash-based prefix caching with MemPool. To realize block aggregation, we
modify several CUDA kernels such as the paged_attention, swap_blocks,
reshape_and_cache.

##  8\. Evaluation

###  8.1. Setup

We describe the physical server, baseline systems, and LLM model used in our
evaluation.

Server. We run all tests on a single NVIDIA DGX H800 server. It has 8
H800-80GB GPUs interconnected by NVLink (400GB/s bandwidth). It has 192-core
Intel Xeon Platinum CPUs @2.4 GHz and 2 TB DRAM. We use Ubuntu 20.04 with
Linux kernel 5.16.7 and CUDA 12.2.

Baseline. We use vLLM-0.4.0 as our baseline for running PD-colocated
instances. MemServe uses the version for building context caching enhanced
disaggregated inference.

Model. We use Llama2-13B with tensor parallel (TP) configured as 2 for all our
tests. We use this model size mainly because it allows us to create four
inference instances within a single server; a larger model would lead to fewer
instances.

Metrics. For end-to-end benchmarking, we report the following metrics: time-
to-first-token (TTFT), job completion time (JCT), and time-per-output-token
(TPOT).

###  8.2. Workloads

Table 7. Workloads Used in Our Work. Type | Workload | Description
---|---|---
Chat | ShareGPT | Chat history with ChatGPT
QA | LooGLE | Long document QA
Agent | ReAct | Agent with acting & reasoning

Figure 7. Workload Statistics. (a) Prompt length. (b) generation length. (c)
The ratio of prompt-len over generated-len. (d) The percentage of the shared
prefix in each workload.

We describe workloads used for our end-to-end tests. We use three
representative workloads as listed in Table 7. (1) ShareGPT (sharegpt, ) is a
real-world dataset containing user-shared ChatGPT conversations. Requests from
the same conversation form a session and share causal dependencies: clients
send a request to the system only after they receive the response to the
conversation’s previous request. We will replace the given response with the
LLM-generated content. (2) LooGLE (li2023loogle, ) is an evaluation benchmark
for LLM long context understanding, containing long documents with QAs.
Similar to ShareGPT, requests constructed from the same document form a
session and share causal dependencies. (3) ReAct (yao2022react, ) is an agent
acting and reasoning framework. We use traces generated from running the ReAct
agent with HotpotQA (yang2018hotpotqa, ) dataset.

Workload Statistics. We study the above workloads in Figure 7. We show the
distribution of prompt and generated token length, their ratio, and the
percentage of shared prefixes representing potential context caching benefits.
Generally, ShareGPT exhibits uniform distribution across all four dimensions.
LooGLE has long prompts, short generation lengths, and a large percentage of
shared prefixes because each request has a long document in its prompt. Since
the document exceeds our model’s context window, we only take the first 1k
tokens of the document and keep the first five associated questions. ReAct
also has long prompts and a large percentage of shared prefixes because each
request has a long two-shot example in its prompt. Unlike LooGLE, requests
from ReAct have relatively long generation lengths because of the long and
thorough reasoning and actions generated from LLMs.

Arrival Pattern. None of the above workloads has arrival patterns. We simulate
a request’s arrival time by sampling it from a Poisson distribution under
different request rates. We maintain the causal dependency for requests
belonging to the same session: a request is only sent to the system after
receiving the response to the session’s previous request.

###  8.3. End-to-End Applications

Figure 8. End-to-End Evalution. The x-axis is the request rate per inference
instance, 1P1D counts as two instances.

We study the benefits of context caching, disaggregated inference, and when
both are combined. We create four different settings: (1) PD denotes PD-
colocated. (2) PD-CC denotes PD-colocated with context caching. (3) 1P1D
denotes disaggregated inference with a single prefill-only and a single
decode-only instance. The numbers can vary. (4) 1P1D-CC denotes 1P1D with
context caching (PC-caching-3). Note that PD-colocated runs vanilla vLLM. The
other three settings are run with MemServe. The request rate is calculated per
instance. Assume a 5 req/s rate, then a 1P1D setup will take 10 req/s. We
ensure an equal number of instances in all tests. Also, we use prompt-tree-
based scheduling and memory aggregation (by-req-agg). All results are in
Figure 8.

ShareGPT. Compared to PD-colocated, disaggregated inference (1P2D over PD)
improves average and P99 JCT by 30% and 42%, respectively. Enhancing
disaggregated inference (e.g., 1P2D) with context caching (e.g., 1P2D-CC)
further improves average and P99 JCT by 17% and 29%, respectively. It also
improves average and P99 TTFT by 58% and 45%, respectively. Since ShareGPT has
the longest generation length of all three workloads, compared to 2P1D, 1P2D
improves JCT because it improves TPOT but at the cost of heavily loaded
prefill instances, hurting TTFT.

LooGLE and ReAct. Both have long prompts and relatively short generation
lengths. For LooGLE, disaggregated inference improves average and P99 JCT by
10.3% and 10.8%, respectively. Context caching further improves average and
P99 JCT by 26.9% and 22.5%, average and P99 TTFT by 56.2% and 45.2%. For
ReAct, disaggregation increases average and P99 JCT by 40.8% and 53.1%.
Caching further enhances these metrics by 26.7% and 21.4%, and average and P99
TTFT by 78.5% and 84.9%.

###  8.4. Microbenchmarks

Figure 9. MemPool API Study. (a) The latency of Memory APIs with varied
numbers of blocks. (a) The latency of key Index APIs with varied cache ratio
and number of blocks.

Figure 10. Caching Study. PD-Colocated. Hash is vanilla vLLM. Radix is an
adapted vLLM with MemPool.

Figure 11. Network and Memory Layout Optimization Study. T is short for
threads. C is short for NCCL communicators. The right figure compares the
performance and HBM usage with varied NCCL buffer sizes. The default is 4 MB.

MemPool API Study. We first study the main MemPool APIs in Figure 9. Without
loss of generality, we show a few key APIs. Memory APIs’ latency increases
linearly with the number of blocks, taking roughly 800 ns per block. For index
APIs, we mainly run insert and match. We vary the number of blocks. A 256
block equals 4K tokens. The latency mostly stays flat with an increasing
cached ratio. It takes at most 0.7 ms to insert a 4K prompt. In all, MemPool
APIs are lightweight and fast.

MemPool Caching Study. We compare vanilla vLLM’s hash-based index with
MemPool’s radix-based index. We run both on a PD-colocated instance with no
cached data. We record the prefill time, which consists of two parts: check
index and model forward. Figure 10 shows that vanilla vLLM’s hash-based prefix
mechanism incurs a huge overhead as the prompt length increases. In all, using
MemPool for basic context caching incurs minimal overhead.

Block Aggregation Study. We study how the proposed memory aggregation helps.
We compare two settings: (1) original discrete memory layout (Original) and
(2) proposed aggregated memory layout (Agg_Block). The test transmits the KV
cache generated from a 2048-token prompt. We tune several key NCCL parameters:
communicator, stream, buffer size, and threads. Figure 11 presents the
results. First, the aggregation method outperforms the vanilla memory layer by
a large margin. Second, a single communicator is enough when the memory block
is large. When the memory block is smaller, multiple communicators are
required for better performance, but as the right figure shows, increasing the
number of communicators consumes extra HBM.

By-Req-Agg Study.

Figure 12. Compare By-Layer, By-Req, and By-Req-Agg.

We run a 1024-prompt-32-decode workload to understand these mechanisms. We
vary the request rate and show results in Figure 12. The proposed by-req-agg
outperforms both by-layer and by-req.

Context Caching Cost Model.

Figure 13. Context Caching Cost Model. All figures have cached-ratio has the
x-axis. Each line represents a different prompt length. All y-axis represent
the TTFT improvement over the no-caching case. (a) studies the prompt-len
factor. (b) studies the batch-size factor. (c) studies the block-size factor.
(d) studies the cached-location factor.

Figure 13 presents the result with several key takeaways. (1) The benefit of
caching improves with a larger cached-ratio. (2) For the same cached-ratio,
longer prompts have higher improvement. (3) Batch size effectively translates
to prompt length. Hence, we need to consider batch size along with cached-
ratio. (4) When the historical KV cache data is located in DRAM, we must swap
it into HBM before using it during prefill. Yet, the benefit of reducing
computation largely offsets the cost of data movement. Regardless of where
data is located, TTFT improves once the cached-ratio exceeds a certain
threshold.

Global Scheduler Study.

Figure 14. Global Scheduler Policy. Share Ratio represents the ratio of the
number of identical requests.

We compare policies as listed in Table 6. We selected 80 sessions from LooGLE,
roughly 250 requests. We propose a share ratio. A share ratio of 2 means
duplicating this set of sessions. While running a 3P1D setup, Figure 14 shows
that compared to intra-session scheduling, prompt-tree-based scheduling
improves P99 TTFT by 59% since it maximizes KV cache reuse.

##  9\. Related Work

Our work is unique in proposing a standalone MemPool module and developing a
holistic serving system MemServe using MemPool APIs.

Disaggregated Inference. Four papers propose the disaggregated inference idea
concurrently within a short 3-month span: Splitwise (patel2023splitwise, ),
TetriServe (tetriserve-2024, ), DistServe (zhong2024distserve, ), and Dejavu
(strati2024dejavu, ). Generally, disaggregating prefill from decode reduces
interference between these two stages and allows each to scale independently
with heterogeneous hardware. More recently, LoongServe (wu2024loongserve, )
takes a step further by enabling dynamic scaling. All prior work builds
disaggregated inference by modifying the inference engine in an ad-hoc manner.
Our work takes a different approach by first abstracting out the MemPool
component and then building disaggregated inference as a use case of MemPool.

Context Caching. Caching reduces recomputation, hence reducing TTFT and
improving throughput. The benefits of context caching are well-studied in
Pensieve (pensieve, ), Cache Gen (liu2023cachegen, ), SGLang (sglang, ), and
Prompt Cache (gim2024promptcache, ). More recently, Google started a
commercial offering of context caching for their Gemini models (gemini-
context-caching, ). All prior work builds context caching in a PD-colocated
setup. Using MemPool APIs, we take a step-by-step approach to building the
first-ever context caching solution atop disaggregated inference.

Scheduling. Scheduling plays a key role in improving serving efficiency. At
the local layer: Orca (yu2022orca, ) proposes iterative-level scheduling to
reduce bubbles. Sarathi (sarathi-arvix23, ; sarathi2, ) proposes chunked-
prefill to overcome suboptimal prefill processing. FastServe(fastserve-
arxiv23, ) utilizes a multi-level priority feedback queue to minimize JCT. At
the global layer: MuxServe (duan2024muxserve, ) formulates a multiplexing
problem and proposes a novel placement algorithm and adaptive batch scheduling
strategy to identify optimal colocations in LLM serving.

Generic Memory Optimization. Many works try to optimize memory usage. For
example, using quantization(rethinking-icml20, ; smoothquant-icml23, ;
Llmint8-arvix23, ; flexgen-pmlr23, ; PTQ-arixv23, ; gptzip-icml23, ; spqr-
arvix22, ; optq-eiclr22, ) to compress the model weights into lower precision,
using paging to reduce fragmentation (vllm-sosp23, ), and low-level algorithm
and kernel optimizations (zeroquant-nips22, ; rammer-osdi20, ; lightseq-
arvix20, ; flashattention-nips22, ; flashattention2-arxiv23, ;
flashdecoding++-2023, ; bytetransformer-ipdps23, ). We refer readers to
(zeng2024cap-llm, ; zhou2024survey-llm, ) for more details.

##  10\. Conclusion

In this paper, we presented MemServe, a novel system designed to enhance the
efficiency of LLM serving by unifying inter-request and intra-request
optimizations. The core of MemServe is a distributed MemPool that manages KV
caches across distributed instances. MemServe builds context caching,
disaggregated inference, and their combo using MemPool APIs. End-to-end
results show MemServe can improve JCT, TTFT, TPOT by a large margin.

## References

  * [1] Amey Agrawal, Nitin Kedia, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S Gulavani, Alexey Tumanov, and Ramachandran Ramjee.  Taming throughput-latency tradeoff in llm inference with sarathi-serve.  arXiv preprint arXiv:2403.02310, 2024.
  * [2] Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S Gulavani, and Ramachandran Ramjee.  Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills.  arXiv preprint arXiv:2308.16369, 2023.
  * [3] Tri Dao.  Flashattention-2: Faster attention with better parallelism and work partitioning.  arXiv preprint arXiv:2307.08691, 2023.
  * [4] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  Advances in Neural Information Processing Systems, 2022.
  * [5] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer.  Llm. int8 (): 8-bit matrix multiplication for transformers at scale.  arXiv preprint arXiv:2208.07339, 2022.
  * [6] Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Saleh Ashkboos, Alexander Borzunov, Torsten Hoefler, and Dan Alistarh.  Spqr: A sparse-quantized representation for near-lossless llm weight compression.  arXiv preprint arXiv:2306.03078, 2023.
  * [7] Jiangfei Duan, Runyu Lu, Haojie Duanmu, Xiuhong Li, Xingcheng Zhang, Dahua Lin, Ion Stoica, and Hao Zhang.  Muxserve: Flexible multiplexing for efficient multiple llm serving.  arXiv preprint arXiv:2404.02015, 2024.
  * [8] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh.  Optq: Accurate quantization for generative pre-trained transformers.  In The Eleventh International Conference on Learning Representations, 2022.
  * [9] In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, and Lin Zhong.  Prompt cache: Modular attention reuse for low-latency inference.  Proceedings of Machine Learning and Systems, 6:325–338, 2024.
  * [10] Google.  Context Caching.  https://ai.google.dev/gemini-api/docs/caching?lang=python.
  * [11] Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, and Yu Wang.  Flashdecoding++: Faster large language model inference on gpus.  arXiv preprint arXiv:2311.01282, 2023.
  * [12] Cunchen Hu, Heyang Huang, Liangliang Xu, Xusheng Chen, Jiang Xu, Shuang Chen, Hao Feng, Chenxi Wang, Sa Wang, Yungang Bao, et al.  Inference without interference: Disaggregate llm inference for mixed downstream workloads.  arXiv preprint arXiv:2401.11181, 2024.
  * [13] Berivan Isik, Hermann Kumbong, Wanyi Ning, Xiaozhe Yao, Sanmi Koyejo, and Ce Zhang.  Gpt-zip: Deep compression of finetuned large language models.  In Workshop on Efficient Systems for Foundation Models@ ICML2023, 2023.
  * [14] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica.  Efficient memory management for large language model serving with pagedattention.  In Proceedings of the 29th Symposium on Operating Systems Principles, 2023.
  * [15] Jiaqi Li, Mengmeng Wang, Zilong Zheng, and Muhan Zhang.  Loogle: Can long-context language models understand long contexts?  arXiv preprint arXiv:2311.04939, 2023.
  * [16] Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez.  Train big, then compress: Rethinking model size for efficient training and inference of transformers.  In International Conference on machine learning, 2020.
  * [17] Bin Lin, Tao Peng, Chen Zhang, Minmin Sun, Lanbo Li, Hanyu Zhao, Wencong Xiao, Qi Xu, Xiafei Qiu, Shen Li, et al.  Infinite-llm: Efficient llm service for long context with distattention and distributed kvcache.  arXiv preprint arXiv:2401.02669, 2024.
  * [18] Yuhan Liu, Hanchen Li, Kuntai Du, Jiayi Yao, Yihua Cheng, Yuyang Huang, Shan Lu, Michael Maire, Henry Hoffmann, Ari Holtzman, et al.  Cachegen: Fast context loading for language model applications.  arXiv preprint arXiv:2310.07240, 2023.
  * [19] Lingxiao Ma, Zhiqiang Xie, Zhi Yang, Jilong Xue, Youshan Miao, Wei Cui, Wenxiang Hu, Fan Yang, Lintao Zhang, and Lidong Zhou.  Rammer: Enabling holistic deep learning compiler optimizations with $\\{$rTasks$\\}$.  In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), 2020.
  * [20] NVIDIA.  NCCL.  https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/overview.html.
  * [21] NVIDIA.  RDMA Verbs.  https://docs.nvidia.com/networking/display/rdmaawareprogrammingv17/rdma+verbs+api.
  * [22] NVIDIA.  TensorRT-LLM.  https://github.com/NVIDIA/TensorRT-LLM.
  * [23] Pratyush Patel, Esha Choukse, Chaojie Zhang, Íñigo Goiri, Aashaka Shah, Saeed Maleki, and Ricardo Bianchini.  Splitwise: Efficient generative llm inference using phase splitting.  arXiv preprint arXiv:2311.18677, 2023.
  * [24] Ramya Prabhu, Ajay Nayak, Jayashree Mohan, Ramachandran Ramjee, and Ashish Panwar.  vattention: Dynamic memory management for serving llms without pagedattention.  arXiv preprint arXiv:2405.04437, 2024.
  * [25] Sharegpt teams.  https://sharegpt.com/.
  * [26] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang.  Flexgen: High-throughput generative inference of large language models with a single gpu.  In International Conference on Machine Learning, 2023.
  * [27] Foteini Strati, Sara Mcallister, Amar Phanishayee, Jakub Tarnawski, and Ana Klimovic.  Déjàvu: Kv-cache streaming for fast, fault-tolerant generative llm serving, 2024.
  * [28] Xiaohui Wang, Ying Xiong, Yang Wei, Mingxuan Wang, and Lei Li.  Lightseq: A high performance inference library for transformers.  arXiv preprint arXiv:2010.13887, 2020.
  * [29] Bingyang Wu, Shengyu Liu, Yinmin Zhong, Peng Sun, Xuanzhe Liu, and Xin Jin.  Loongserve: Efficiently serving long-context large language models with elastic sequence parallelism.  arXiv preprint arXiv:2404.09526, 2024.
  * [30] Bingyang Wu, Yinmin Zhong, Zili Zhang, Gang Huang, Xuanzhe Liu, and Xin Jin.  Fast distributed inference serving for large language models.  arXiv preprint arXiv:2305.05920, 2023.
  * [31] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han.  Smoothquant: Accurate and efficient post-training quantization for large language models.  In International Conference on Machine Learning, 2023.
  * [32] Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning.  Hotpotqa: A dataset for diverse, explainable multi-hop question answering.  arXiv preprint arXiv:1809.09600, 2018.
  * [33] Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao.  React: Synergizing reasoning and acting in language models.  arXiv preprint arXiv:2210.03629, 2022.
  * [34] Zhewei Yao, Cheng Li, Xiaoxia Wu, Stephen Youn, and Yuxiong He.  A comprehensive study on post-training quantization for large language models.  arXiv preprint arXiv:2303.08302, 2023.
  * [35] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He.  Zeroquant: Efficient and affordable post-training quantization for large-scale transformers.  Advances in Neural Information Processing Systems, 2022.
  * [36] Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun.  Orca: A distributed serving system for $\\{$Transformer-Based$\\}$ generative models.  In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 2022.
  * [37] Lingfan Yu and Jinyang Li.  Stateful large language model serving with pensieve.  arXiv preprint arXiv:2312.05516, 2023.
  * [38] Pai Zeng, Zhenyu Ning, Jieru Zhao, Weihao Cui, Mengwei Xu, Liwei Guo, Xusheng Chen, and Yizhou Shan.  The cap principle for llm serving.  arXiv preprint arXiv:2405.11299, 2024.
  * [39] Yujia Zhai, Chengquan Jiang, Leyuan Wang, Xiaoying Jia, Shang Zhang, Zizhong Chen, Xin Liu, and Yibo Zhu.  Bytetransformer: A high-performance transformer boosted for variable-length inputs.  In 2023 IEEE International Parallel and Distributed Processing Symposium (IPDPS), 2023.
  * [40] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Jeff Huang, Chuyue Sun, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al.  Efficiently programming large language models using sglang.  arXiv preprint arXiv:2312.07104, 2023.
  * [41] Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang.  Distserve: Disaggregating prefill and decoding for goodput-optimized large language model serving, 2024.
  * [42] Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, et al.  A survey on efficient inference for large language models.  arXiv preprint arXiv:2404.14294, 2024.