diff --git a/level_zero/core/source/cmdqueue/cmdqueue_hw_base.inl b/level_zero/core/source/cmdqueue/cmdqueue_hw_base.inl
index 2bd9fd671b..de900aed83 100644
--- a/level_zero/core/source/cmdqueue/cmdqueue_hw_base.inl
+++ b/level_zero/core/source/cmdqueue/cmdqueue_hw_base.inl
@@ -40,7 +40,7 @@ void CommandQueueHw<gfxCoreFamily>::programStateBaseAddress(uint64_t gsba, bool
     NEO::EncodeWA<GfxFamily>::addPipeControlBeforeStateBaseAddress(commandStream, hwInfo, isRcs);
     NEO::EncodeWA<GfxFamily>::encodeAdditionalPipelineSelect(commandStream, {}, true, hwInfo, isRcs);
 
-    auto sbaCmdBuf = static_cast<STATE_BASE_ADDRESS *>(NEO::StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStream));
+    auto sbaCmdBuf = NEO::StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStream);
     STATE_BASE_ADDRESS sbaCmd;
 
     bool useGlobalSshAndDsh = NEO::ApiSpecificConfig::getBindlessConfiguration();
diff --git a/level_zero/core/source/cmdqueue/cmdqueue_xe_hp_core_and_later.inl b/level_zero/core/source/cmdqueue/cmdqueue_xe_hp_core_and_later.inl
index d55acb406f..5d753c0f53 100644
--- a/level_zero/core/source/cmdqueue/cmdqueue_xe_hp_core_and_later.inl
+++ b/level_zero/core/source/cmdqueue/cmdqueue_xe_hp_core_and_later.inl
@@ -40,7 +40,7 @@ void CommandQueueHw<gfxCoreFamily>::programStateBaseAddress(uint64_t gsba, bool
         bool isRcs = this->getCsr()->isRcs();
 
         NEO::EncodeWA<GfxFamily>::addPipeControlBeforeStateBaseAddress(commandStream, hwInfo, isRcs);
-        auto sbaCmdBuf = static_cast<STATE_BASE_ADDRESS *>(NEO::StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStream));
+        auto sbaCmdBuf = NEO::StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStream);
 
         STATE_BASE_ADDRESS sbaCmd;
         NEO::StateBaseAddressHelperArgs<GfxFamily> args = {
@@ -68,7 +68,7 @@ void CommandQueueHw<gfxCoreFamily>::programStateBaseAddress(uint64_t gsba, bool
 
         auto &hwInfoConfig = *NEO::HwInfoConfig::get(hwInfo.platform.eProductFamily);
         if (hwInfoConfig.isAdditionalStateBaseAddressWARequired(hwInfo)) {
-            sbaCmdBuf = static_cast<STATE_BASE_ADDRESS *>(commandStream.getSpace(sizeof(STATE_BASE_ADDRESS)));
+            sbaCmdBuf = NEO::StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStream);
             *sbaCmdBuf = sbaCmd;
         }
 
diff --git a/shared/source/command_container/command_encoder_bdw_and_later.inl b/shared/source/command_container/command_encoder_bdw_and_later.inl
index ce96a5b543..a8b2be4824 100644
--- a/shared/source/command_container/command_encoder_bdw_and_later.inl
+++ b/shared/source/command_container/command_encoder_bdw_and_later.inl
@@ -393,7 +393,7 @@ void EncodeStateBaseAddress<Family>::encode(CommandContainer &container, STATE_B
 
     StateBaseAddressHelper<Family>::programStateBaseAddress(args);
 
-    auto cmdSpace = reinterpret_cast<STATE_BASE_ADDRESS *>(container.getCommandStream()->getSpace(sizeof(STATE_BASE_ADDRESS)));
+    auto cmdSpace = StateBaseAddressHelper<Family>::getSpaceForSbaCmd(*container.getCommandStream());
     *cmdSpace = sbaCmd;
 
     EncodeWA<Family>::encodeAdditionalPipelineSelect(*container.getCommandStream(), {}, false, hwInfo, isRcs);
diff --git a/shared/source/command_container/command_encoder_xehp_and_later.inl b/shared/source/command_container/command_encoder_xehp_and_later.inl
index e649685f7d..474c820613 100644
--- a/shared/source/command_container/command_encoder_xehp_and_later.inl
+++ b/shared/source/command_container/command_encoder_xehp_and_later.inl
@@ -500,13 +500,13 @@ void EncodeStateBaseAddress<Family>::encode(CommandContainer &container, STATE_B
 
     StateBaseAddressHelper<Family>::programStateBaseAddress(args);
 
-    auto cmdSpace = reinterpret_cast<STATE_BASE_ADDRESS *>(container.getCommandStream()->getSpace(sizeof(STATE_BASE_ADDRESS)));
+    auto cmdSpace = StateBaseAddressHelper<Family>::getSpaceForSbaCmd(*container.getCommandStream());
     *cmdSpace = sbaCmd;
 
     auto &hwInfo = container.getDevice()->getHardwareInfo();
     auto &hwInfoConfig = *HwInfoConfig::get(hwInfo.platform.eProductFamily);
     if (hwInfoConfig.isAdditionalStateBaseAddressWARequired(hwInfo)) {
-        cmdSpace = reinterpret_cast<STATE_BASE_ADDRESS *>(container.getCommandStream()->getSpace(sizeof(STATE_BASE_ADDRESS)));
+        cmdSpace = StateBaseAddressHelper<Family>::getSpaceForSbaCmd(*container.getCommandStream());
         *cmdSpace = sbaCmd;
     }
 
diff --git a/shared/source/command_stream/command_stream_receiver_hw_base.inl b/shared/source/command_stream/command_stream_receiver_hw_base.inl
index 06ad7fd608..f5db05d430 100644
--- a/shared/source/command_stream/command_stream_receiver_hw_base.inl
+++ b/shared/source/command_stream/command_stream_receiver_hw_base.inl
@@ -403,7 +403,7 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
         }
 
         auto stateBaseAddressCmdOffset = commandStreamCSR.getUsed();
-        auto stateBaseAddressCmdBuffer = static_cast<STATE_BASE_ADDRESS *>(StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStreamCSR));
+        auto stateBaseAddressCmdBuffer = StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(commandStreamCSR);
         auto instructionHeapBaseAddress = getMemoryManager()->getInternalHeapBaseAddress(rootDeviceIndex, getMemoryManager()->isLocalMemoryUsedForIsa(rootDeviceIndex));
         uint64_t indirectObjectStateBaseAddress = getMemoryManager()->getInternalHeapBaseAddress(rootDeviceIndex, ioh->getGraphicsAllocation()->isAllocatedInLocalMemoryPool());
 
diff --git a/shared/source/helpers/state_base_address.h b/shared/source/helpers/state_base_address.h
index 228221bf20..2d534bbf6b 100644
--- a/shared/source/helpers/state_base_address.h
+++ b/shared/source/helpers/state_base_address.h
@@ -49,7 +49,7 @@ template <typename GfxFamily>
 struct StateBaseAddressHelper {
     using STATE_BASE_ADDRESS = typename GfxFamily::STATE_BASE_ADDRESS;
 
-    static void *getSpaceForSbaCmd(LinearStream &cmdStream);
+    static STATE_BASE_ADDRESS *getSpaceForSbaCmd(LinearStream &cmdStream);
 
     static void programStateBaseAddress(StateBaseAddressHelperArgs<GfxFamily> &args);
 
diff --git a/shared/source/helpers/state_base_address_base.inl b/shared/source/helpers/state_base_address_base.inl
index 65b39ef192..c4f7676713 100644
--- a/shared/source/helpers/state_base_address_base.inl
+++ b/shared/source/helpers/state_base_address_base.inl
@@ -85,8 +85,8 @@ void StateBaseAddressHelper<GfxFamily>::programStateBaseAddress(
 }
 
 template <typename GfxFamily>
-void *StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(LinearStream &cmdStream) {
-    return cmdStream.getSpace(sizeof(STATE_BASE_ADDRESS));
+typename GfxFamily::STATE_BASE_ADDRESS *StateBaseAddressHelper<GfxFamily>::getSpaceForSbaCmd(LinearStream &cmdStream) {
+    return cmdStream.getSpaceForCmd<typename GfxFamily::STATE_BASE_ADDRESS>();
 }
 
 template <typename GfxFamily>
diff --git a/shared/source/xe_hpg_core/command_stream_receiver_hw_xe_hpg_core.cpp b/shared/source/xe_hpg_core/command_stream_receiver_hw_xe_hpg_core.cpp
index 14a1475f40..02ee544155 100644
--- a/shared/source/xe_hpg_core/command_stream_receiver_hw_xe_hpg_core.cpp
+++ b/shared/source/xe_hpg_core/command_stream_receiver_hw_xe_hpg_core.cpp
@@ -48,8 +48,8 @@ void CommandStreamReceiverHw<Family>::programAdditionalStateBaseAddress(LinearSt
     auto &hwInfo = *device.getRootDeviceEnvironment().getHardwareInfo();
     auto &hwInfoConfig = *HwInfoConfig::get(hwInfo.platform.eProductFamily);
     if (hwInfoConfig.isAdditionalStateBaseAddressWARequired(hwInfo)) {
-        auto pCmd = static_cast<STATE_BASE_ADDRESS *>(csr.getSpace(sizeof(STATE_BASE_ADDRESS)));
-        *pCmd = cmd;
+        auto cmdSpace = StateBaseAddressHelper<Family>::getSpaceForSbaCmd(csr);
+        *cmdSpace = cmd;
     }
 }